feat: flattern nested arithmetic (#594)

pkg/hir/lower.go

@@ -299,15 +299,35 @@ func extractBodies(es []Expr, schema *mir.Schema) []mir.Expr {
 func expand(e Expr, schema sc.Schema) []Expr {
 	if p, ok := e.(*Add); ok {
 		return expandWithNaryConstructor(p.Args, func(nargs []Expr) Expr {
-			return &Add{Args: nargs}
+			var args []Expr
+			// Flatten nested sums
+			for _, e := range nargs {
+				if a, ok := e.(*Add); ok {
+					args = append(args, a.Args...)
+				} else {
+					args = append(args, e)
+				}
+			}
+			// Done
+			return &Add{Args: args}
 		}, schema)
 	} else if _, ok := e.(*Constant); ok {
 		return []Expr{e}
 	} else if _, ok := e.(*ColumnAccess); ok {
 		return []Expr{e}
 	} else if p, ok := e.(*Mul); ok {
 		return expandWithNaryConstructor(p.Args, func(nargs []Expr) Expr {
-			return &Mul{Args: nargs}
+			var args []Expr
+			// Flatten nested products
+			for _, e := range nargs {
+				if a, ok := e.(*Mul); ok {
+					args = append(args, a.Args...)
+				} else {
+					args = append(args, e)
+				}
+			}
+			// Done
+			return &Mul{Args: args}
 		}, schema)
 	} else if p, ok := e.(*List); ok {
 		ees := make([]Expr, 0)