Match expressions.

powdr-labs · Sep 24, 2024 · 26930b1 · 26930b1
1 parent 86b0a2d
commit 26930b1
Showing 1 changed file with 47 additions and 4 deletions.
diff --git a/jit-compiler/src/codegen.rs b/jit-compiler/src/codegen.rs
@@ -7,7 +7,8 @@ use powdr_ast::{
         display::{format_type_args, quote},
         types::{ArrayType, FunctionType, Type, TypeScheme},
         ArrayLiteral, BinaryOperation, BinaryOperator, BlockExpression, FunctionCall, IfExpression,
-        IndexAccess, LambdaExpression, Number, StatementInsideBlock, UnaryOperation,
+        IndexAccess, LambdaExpression, MatchArm, MatchExpression, Number, Pattern,
+        StatementInsideBlock, UnaryOperation,
     },
 };
 use powdr_number::FieldElement;
@@ -268,10 +269,50 @@ impl<'a, T: FieldElement> CodeGenerator<'a, T> {
                         .unwrap_or_default()
                 )
             }
+            Expression::MatchExpression(_, MatchExpression { scrutinee, arms }) => {
+                format!(
+                    "match {} {{\n{}\n}}",
+                    self.format_expr(scrutinee)?,
+                    arms.iter()
+                        .map(|MatchArm { pattern, value }| {
+                            Ok(format!(
+                                "{} => {},",
+                                self.format_pattern(pattern),
+                                self.format_expr(value)?,
+                            ))
+                        })
+                        .collect::<Result<Vec<_>, String>>()?
+                        .join("\n")
+                )
+            }
             _ => return Err(format!("Implement {e}")),
         })
     }
 
+    fn format_pattern(&mut self, pattern: &Pattern) -> String {
+        match pattern {
+            Pattern::CatchAll(_) => "_".to_string(),
+            Pattern::Ellipsis(_) => "..".to_string(),
+            Pattern::Number(_, n) => n.to_string(),
+            Pattern::String(_, s) => quote(s),
+            Pattern::Tuple(_, items) => format!(
+                "({})",
+                items.iter().map(|p| self.format_pattern(p)).join(", ")
+            ),
+            Pattern::Array(_, items) => format!(
+                "[{}]",
+                items.iter().map(|p| self.format_pattern(p)).join(", ")
+            ),
+            Pattern::Variable(_, var) => var.clone(),
+            Pattern::Enum(_, name, None) => escape_symbol(&name.to_string()),
+            Pattern::Enum(_, name, Some(fields)) => format!(
+                "{}({})",
+                escape_symbol(&name.to_string()),
+                fields.iter().map(|p| self.format_pattern(p)).join(", ")
+            ),
+        }
+    }
+
     fn format_statement(&mut self, s: &StatementInsideBlock<Expression>) -> Result<String, String> {
         Err(format!("Implement {s}"))
     }
@@ -294,10 +335,12 @@ fn map_type(ty: &Type) -> String {
         Type::Function(ft) => todo!("Type {ft}"),
         Type::TypeVar(tv) => tv.to_string(),
         Type::NamedType(path, type_args) => {
-            if type_args.is_some() {
-                unimplemented!()
+            let name = escape_symbol(&path.to_string());
+            if let Some(type_args) = type_args {
+                format!("{name}::<{}>", type_args.iter().map(map_type).join(", "))
+            } else {
+                name
             }
-            escape_symbol(&path.to_string())
         }
         Type::Col | Type::Inter => unreachable!(),
     }