datatype*

src/ast/desugar.rs

@@ -364,6 +364,49 @@ pub(crate) fn desugar_command(
         Command::Input { span, name, file } => {
             vec![NCommand::Input { span, name, file }]
         }
+        Command::Datatypes { span: _, datatypes } => {
+            let mut res = vec![];
+            for datatype in datatypes.iter() {
+                let span = datatype.0.clone();
+                let name = datatype.1;
+                if datatype.2.is_ok() {
+                    res.push(NCommand::Sort(span, name, None));
+                }
+            }
+            let (variants_vec, sorts): (Vec<_>, Vec<_>) = datatypes
+                .into_iter()
+                .partition(|datatype| datatype.2.is_ok());
+
+            for sort in sorts {
+                let span = sort.0.clone();
+                let name = sort.1;
+                let constructor = sort.2.unwrap_err();
+                res.push(NCommand::Sort(span, name, Some(constructor)));
+            }
+
+            for variants in variants_vec {
+                let datatype = variants.1;
+                let variants = variants.2.unwrap();
+                for variant in variants {
+                    res.push(NCommand::Function(FunctionDecl {
+                        name: variant.name,
+                        schema: Schema {
+                            input: variant.types,
+                            output: datatype,
+                        },
+                        merge: None,
+                        merge_action: Actions::default(),
+                        default: None,
+                        cost: variant.cost,
+                        unextractable: false,
+                        ignore_viz: false,
+                        span: variant.span,
+                    }));
+                }
+            }
+
+            res
+        }
     };
 
     Ok(res)

src/ast/mod.rs

@@ -435,6 +435,8 @@ pub type Command = GenericCommand<Symbol, Symbol>;
 
 pub type Subsume = bool;
 
+pub type Subdatatypes = Result<Vec<Variant>, (Symbol, Vec<Expr>)>;
+
 /// A [`Command`] is the top-level construct in egglog.
 /// It includes defining rules, declaring functions,
 /// adding to tables, and running rules (via a [`Schedule`]).
@@ -488,6 +490,10 @@ where
         name: Symbol,
         variants: Vec<Variant>,
     },
+    Datatypes {
+        span: Span,
+        datatypes: Vec<(Span, Symbol, Subdatatypes)>,
+    },
     /// Create a new user-defined sort, which can then
     /// be used in new [`Command::Function`] declarations.
     /// The [`Command::Datatype`] command desugars directly to this command, with one [`Command::Function`]
@@ -503,11 +509,7 @@ where
     /// ```
     ///
     /// Now `MathVec` can be used as an input or output sort.
-    Sort(
-        Span,
-        Symbol,
-        Option<(Symbol, Vec<GenericExpr<Symbol, Symbol>>)>,
-    ),
+    Sort(Span, Symbol, Option<(Symbol, Vec<Expr>)>),
     /// Declare an egglog function, which is a database table with a
     /// a functional dependency (also called a primary key) on its inputs to one output.
     ///
@@ -878,6 +880,16 @@ where
                 expr,
                 schedule,
             } => list!("simplify", schedule, expr),
+            GenericCommand::Datatypes { span: _, datatypes } => {
+                let datatypes: Vec<_> = datatypes
+                    .iter()
+                    .map(|(_, name, variants)| match variants {
+                        Ok(variants) => list!(name, ++ variants),
+                        Err((head, args)) => list!("sort", name, list!(head, ++ args)),
+                    })
+                    .collect();
+                list!("datatypes", ++ datatypes)
+            }
         }
     }
 }

src/ast/parse.lalrpop

@@ -44,11 +44,19 @@ Comma<T>: Vec<T> = {
     }
 };
 
+RecDatatype: (Span, Symbol, Result<Vec<Variant>, (Symbol, Vec<Expr>)>) = {
+    <lo:LParen> <name:Ident> <variants:(Variant)*> <hi:RParen> => (Span(srcfile.clone(), lo, hi), name, Ok(variants)),
+    <lo:LParen> "sort" <name:Ident> LParen <head:Ident> <exprs:(Expr)*> RParen <hi:RParen> => (Span(srcfile.clone(), lo, hi), name, Err((head, exprs))),
+}
+
 Command: Command = {
     LParen "set-option" <name:Ident> <value:Expr> RParen => Command::SetOption { name, value },
     <lo:LParen> "datatype" <name:Ident> <variants:(Variant)*> <hi:RParen> => Command::Datatype { span: Span(srcfile.clone(), lo, hi), name, variants },
     <lo:LParen> "sort" <name:Ident> LParen <head:Ident> <tail:(Expr)*> RParen <hi:RParen> => Command::Sort (Span(srcfile.clone(), lo, hi), name, Some((head, tail))),
     <lo:LParen> "sort" <name:Ident> <hi:RParen> => Command::Sort (Span(srcfile.clone(), lo, hi), name, None),
+    <lo:LParen> "datatype*"
+        <datatypes:RecDatatype*>
+    <hi:RParen> => Command::Datatypes { span: Span(srcfile.clone(), lo, hi), datatypes },
     <lo:LParen> "function" <name:Ident> <schema:Schema> <cost:Cost>
         <unextractable:(":unextractable")?>
         <merge_action:(":on_merge" <List<Action>>)?>

tests/datatypes.egg

@@ -0,0 +1,11 @@
+(datatype*
+    (Math
+        (Add Math Math)
+        (Sum MathVec)
+        (B Bool))
+    (sort MathVec (Vec Math))
+    (Bool
+        (True)
+        (False)))
+
+(let expr (Add (Sum (vec-of (B (True)) (B (False)))) (B (True))))