Add if statement support

src/ast.rs

@@ -239,6 +239,8 @@ pub enum SingleExpressionInner {
     Call(Call),
     /// Match expression.
     Match(Match),
+    /// If expression.
+    If(If),
 }
 
 /// Call of a user-defined or of a builtin function.
@@ -403,6 +405,38 @@ impl MatchArm {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct If {
+    scrutinee: Arc<Expression>,
+    then_arm: Arc<Expression>,
+    else_arm: Arc<Expression>,
+    span: Span,
+}
+
+impl If {
+    /// Access the expression who's output is deconstructed in the `if`.
+    pub fn scrutinee(&self) -> &Expression {
+        &self.scrutinee
+    }
+
+    /// Access the branch that handles the `true` portion of the `if`.
+    pub fn then_arm(&self) -> &Expression {
+        &self.then_arm
+    }
+
+    /// Access the branch that handles the `false` or `else` portion of the `if`.
+    pub fn else_arm(&self) -> &Expression {
+        &self.else_arm
+    }
+
+    /// Access the span of the if statement.
+    pub fn span(&self) -> &Span {
+        &self.span
+    }
+}
+
+impl_eq_hash!(If; scrutinee, then_arm, else_arm);
+
 /// Item when analyzing modules.
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum ModuleItem {
@@ -462,6 +496,7 @@ pub enum ExprTree<'a> {
     Single(&'a SingleExpression),
     Call(&'a Call),
     Match(&'a Match),
+    If(&'a If),
 }
 
 impl TreeLike for ExprTree<'_> {
@@ -502,13 +537,19 @@ impl TreeLike for ExprTree<'_> {
                 }
                 S::Call(call) => Tree::Unary(Self::Call(call)),
                 S::Match(match_) => Tree::Unary(Self::Match(match_)),
+                S::If(if_) => Tree::Unary(Self::If(if_)),
             },
             Self::Call(call) => Tree::Nary(call.args().iter().map(Self::Expression).collect()),
             Self::Match(match_) => Tree::Nary(Arc::new([
                 Self::Expression(match_.scrutinee()),
                 Self::Expression(match_.left().expression()),
                 Self::Expression(match_.right().expression()),
             ])),
+            Self::If(if_) => Tree::Nary(Arc::new([
+                Self::Expression(if_.scrutinee()),
+                Self::Expression(if_.then_arm()),
+                Self::Expression(if_.else_arm()),
+            ])),
         }
     }
 }
@@ -1059,6 +1100,9 @@ impl AbstractSyntaxTree for SingleExpression {
             parse::SingleExpressionInner::Match(match_) => {
                 Match::analyze(match_, ty, scope).map(SingleExpressionInner::Match)?
             }
+            parse::SingleExpressionInner::If(if_) => {
+                If::analyze(if_, ty, scope).map(SingleExpressionInner::If)?
+            }
         };
 
         Ok(Self {
@@ -1426,6 +1470,28 @@ impl AbstractSyntaxTree for Match {
     }
 }
 
+impl AbstractSyntaxTree for If {
+    type From = parse::If;
+
+    fn analyze(from: &Self::From, ty: &ResolvedType, scope: &mut Scope) -> Result<Self, RichError> {
+        let scrutinee =
+            Expression::analyze(from.scrutinee(), &ResolvedType::boolean(), scope).map(Arc::new)?;
+        scope.push_scope();
+        let ast_then = Expression::analyze(from.then_arm(), ty, scope).map(Arc::new)?;
+        scope.pop_scope();
+        scope.push_scope();
+        let ast_else = Expression::analyze(from.else_arm(), ty, scope).map(Arc::new)?;
+        scope.pop_scope();
+
+        Ok(Self {
+            scrutinee,
+            then_arm: ast_then,
+            else_arm: ast_else,
+            span: *from.as_ref(),
+        })
+    }
+}
+
 fn analyze_named_module(
     name: ModuleName,
     from: &parse::ModuleProgram,
@@ -1559,6 +1625,12 @@ impl AsRef<Span> for Match {
     }
 }
 
+impl AsRef<Span> for If {
+    fn as_ref(&self) -> &Span {
+        &self.span
+    }
+}
+
 impl AsRef<Span> for Module {
     fn as_ref(&self) -> &Span {
         &self.span

src/compile/mod.rs

@@ -12,7 +12,7 @@ use simplicity::{types, Cmr, FailEntropy};
 use self::builtins::array_fold;
 use crate::array::{BTreeSlice, Partition};
 use crate::ast::{
-    Call, CallName, Expression, ExpressionInner, Match, Program, SingleExpression,
+    Call, CallName, Expression, ExpressionInner, If, Match, Program, SingleExpression,
     SingleExpressionInner, Statement,
 };
 use crate::debug::CallTracker;
@@ -355,6 +355,7 @@ impl SingleExpression {
             }
             SingleExpressionInner::Call(call) => call.compile(scope)?,
             SingleExpressionInner::Match(match_) => match_.compile(scope)?,
+            SingleExpressionInner::If(if_) => if_.compile(scope)?,
         };
 
         scope
@@ -680,3 +681,22 @@ impl Match {
         input.comp(&output).with_span(self)
     }
 }
+
+impl If {
+    fn compile<'brand>(
+        &self,
+        scope: &mut Scope<'brand>,
+    ) -> Result<PairBuilder<ProgNode<'brand>>, RichError> {
+        scope.push_scope();
+        let then_arm = self.then_arm().compile(scope)?;
+        scope.pop_scope();
+        scope.push_scope();
+        let else_arm = self.else_arm().compile(scope)?;
+        scope.pop_scope();
+
+        let scrutinee = self.scrutinee().compile(scope)?;
+        let input = scrutinee.pair(PairBuilder::iden(scope.ctx()));
+        let output = ProgNode::case(then_arm.as_ref(), else_arm.as_ref()).with_span(self)?;
+        input.comp(&output).with_span(self)
+    }
+}

src/lexer.rs

@@ -17,21 +17,37 @@ pub enum Token<'src> {
     Mod,
     Const,
     Match,
+    If,
+    Else,
 
     // Control symbols
+    /// `->``
     Arrow,
+    /// `:`
     Colon,
+    /// `;`
     Semi,
+    /// `,`
     Comma,
+    /// `=`
     Eq,
+    /// `=>`
     FatArrow,
+    /// `(`
     LParen,
+    /// `)`
     RParen,
+    /// `[`
     LBracket,
+    /// `]`
     RBracket,
+    /// `{`
     LBrace,
+    /// `}`
     RBrace,
+    /// `<`
     LAngle,
+    /// `>`
     RAngle,
 
     // Number literals
@@ -69,6 +85,8 @@ impl<'src> fmt::Display for Token<'src> {
             Token::Mod => write!(f, "mod"),
             Token::Const => write!(f, "const"),
             Token::Match => write!(f, "match"),
+            Token::If => write!(f, "if"),
+            Token::Else => write!(f, "else"),
 
             Token::Arrow => write!(f, "->"),
             Token::Colon => write!(f, ":"),
@@ -140,6 +158,8 @@ pub fn lexer<'src>(
         "mod" => Token::Mod,
         "const" => Token::Const,
         "match" => Token::Match,
+        "if" => Token::If,
+        "else" => Token::Else, // TODO: Else is never parsed.
         "true" => Token::Bool(true),
         "false" => Token::Bool(false),
         _ => Token::Ident(s),
@@ -251,6 +271,88 @@ mod tests {
         (tokens, errors)
     }
 
+    /// Helper function to get the variant name of a token
+    fn variant_name(token: &Token) -> &'static str {
+        match token {
+            Token::Fn => "Fn",
+            Token::Let => "Let",
+            Token::Type => "Type",
+            Token::Mod => "Mod",
+            Token::Const => "Const",
+            Token::Match => "Match",
+            Token::If => "If",
+            Token::Else => "Else",
+            Token::Arrow => "Arrow",
+            Token::Colon => "Colon",
+            Token::Semi => "Semi",
+            Token::Comma => "Comma",
+            Token::Eq => "Eq",
+            Token::FatArrow => "FatArrow",
+            Token::LParen => "LParen",
+            Token::RParen => "RParen",
+            Token::LBracket => "LBracket",
+            Token::RBracket => "RBracket",
+            Token::LBrace => "LBrace",
+            Token::RBrace => "RBrace",
+            Token::LAngle => "LAngle",
+            Token::RAngle => "RAngle",
+            Token::DecLiteral(_) => "DecLiteral",
+            Token::HexLiteral(_) => "HexLiteral",
+            Token::BinLiteral(_) => "BinLiteral",
+            Token::Bool(_) => "Bool",
+            Token::Ident(_) => "Ident",
+            Token::Jet(_) => "Jet",
+            Token::Witness(_) => "Witness",
+            Token::Param(_) => "Param",
+            Token::Macro(_) => "Macro",
+            Token::Comment => "Comment",
+            Token::BlockComment => "BlockComment",
+        }
+    }
+
+    /// Macro to assert that a sequence of tokens matches the expected variant types
+    macro_rules! assert_tokens_match {
+        ($tokens:expr, $($expected:ident),* $(,)?) => {
+            {
+                let tokens = $tokens.as_ref().expect("Expected Some tokens");
+                let expected_variants = vec![$( stringify!($expected) ),*];
+
+                assert_eq!(
+                    tokens.len(),
+                    expected_variants.len(),
+                    "Expected {} tokens, got {}.\nTokens: {:?}",
+                    expected_variants.len(),
+                    tokens.len(),
+                    tokens
+                );
+
+                for (idx, (token, expected_variant)) in tokens.iter().zip(expected_variants.iter()).enumerate() {
+                    let actual_variant = variant_name(token);
+                    assert_eq!(
+                        actual_variant,
+                        *expected_variant,
+                        "Token at index {} does not match: expected {}, got {} (token: {:?})",
+                        idx,
+                        expected_variant,
+                        actual_variant,
+                        token
+                    );
+                }
+            }
+        };
+    }
+
+    #[test]
+    fn test_if_statement() {
+        let input = "if true {0} else {1};";
+        let (tokens, errors) = lex(input);
+        assert!(errors.is_empty(), "Expected no errors, found: {:?}", errors);
+
+        assert_tokens_match!(
+            tokens, If, Bool, LBrace, DecLiteral, RBrace, Else, LBrace, DecLiteral, RBrace, Semi,
+        );
+    }
+
     #[test]
     fn test_block_comment_simple() {
         let input = "/* hello world */";