Make lambda not break in some cases

more_math/Parser/MathExpr.g4

@@ -49,7 +49,10 @@ breakStmt: BREAK SEMICOLON;
 continueStmt: CONTINUE SEMICOLON;
 returnStmt: RETURN expr? SEMICOLON;
 
-expr: ternaryExpr | compExpr;
+// A lambda binds loosest, so its body takes the whole following expression, ternaries included.
+expr: lambdaExpr | ternaryExpr | compExpr;
+
+lambdaExpr: LPAREN paramList RPAREN ARROW (expr | block) # LambdaExp;
 
 ternaryExpr: compExpr QUESTION expr COLON expr # TernaryExp;
 
@@ -117,7 +120,6 @@ atom:
 	| NUMBER						# NumberExp
 	| CONSTANT						# ConstantExp
 	| STRING						# StringExp
-	| LPAREN paramList RPAREN ARROW (expr | block)	# LambdaExp
 	| LPAREN expr RPAREN			# ParenExp
 	| LPAREN expr RPAREN			# ParenExp
 	| PIPE expr PIPE				# AbsExp

more_math/Parser/UnifiedMathVisitor.py

@@ -3161,8 +3161,11 @@ class UnifiedMathVisitor(MathExprVisitor):
                 if res is not None:
                     last_res = res
         except ReturnSignal as signal:
-            return signal.value
+            last_res = signal.value
 
+        if isinstance(last_res, LambdaFunction):
+            final = ctx.expr()
+            raise ValueError(f"{final.start.line}:{final.start.column}: the script ends with a lambda, but it has to end with a value; call it instead, e.g. f(V0)")
         return last_res
 
     def visitExprStatement(self, ctx):

more_math/Parser/grammer/MathExprListener.py

@@ -215,6 +215,15 @@ class MathExprListener(ParseTreeListener):
         pass
 
 
+    # Enter a parse tree produced by MathExprParser#LambdaExp.
+    def enterLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
+        pass
+
+    # Exit a parse tree produced by MathExprParser#LambdaExp.
+    def exitLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
+        pass
+
+
     # Enter a parse tree produced by MathExprParser#TernaryExp.
     def enterTernaryExp(self, ctx:MathExprParser.TernaryExpContext):
         pass
@@ -638,15 +647,6 @@ class MathExprListener(ParseTreeListener):
         pass
 
 
-    # Enter a parse tree produced by MathExprParser#LambdaExp.
-    def enterLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
-        pass
-
-    # Exit a parse tree produced by MathExprParser#LambdaExp.
-    def exitLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
-        pass
-
-
     # Enter a parse tree produced by MathExprParser#ParenExp.
     def enterParenExp(self, ctx:MathExprParser.ParenExpContext):
         pass

more_math/Parser/grammer/MathExprVisitor.py

@@ -124,6 +124,11 @@ class MathExprVisitor(ParseTreeVisitor):
         return self.visitChildren(ctx)
 
 
+    # Visit a parse tree produced by MathExprParser#LambdaExp.
+    def visitLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
+        return self.visitChildren(ctx)
+
+
     # Visit a parse tree produced by MathExprParser#TernaryExp.
     def visitTernaryExp(self, ctx:MathExprParser.TernaryExpContext):
         return self.visitChildren(ctx)
@@ -359,11 +364,6 @@ class MathExprVisitor(ParseTreeVisitor):
         return self.visitChildren(ctx)
 
 
-    # Visit a parse tree produced by MathExprParser#LambdaExp.
-    def visitLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
-        return self.visitChildren(ctx)
-
-
     # Visit a parse tree produced by MathExprParser#ParenExp.
     def visitParenExp(self, ctx:MathExprParser.ParenExpContext):
         return self.visitChildren(ctx)

more_math/Parser/legacy/MathExprListener.py

@@ -215,6 +215,15 @@ class MathExprListener(ParseTreeListener):
         pass
 
 
+    # Enter a parse tree produced by MathExprParser#LambdaExp.
+    def enterLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
+        pass
+
+    # Exit a parse tree produced by MathExprParser#LambdaExp.
+    def exitLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
+        pass
+
+
     # Enter a parse tree produced by MathExprParser#TernaryExp.
     def enterTernaryExp(self, ctx:MathExprParser.TernaryExpContext):
         pass
@@ -638,15 +647,6 @@ class MathExprListener(ParseTreeListener):
         pass
 
 
-    # Enter a parse tree produced by MathExprParser#LambdaExp.
-    def enterLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
-        pass
-
-    # Exit a parse tree produced by MathExprParser#LambdaExp.
-    def exitLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
-        pass
-
-
     # Enter a parse tree produced by MathExprParser#ParenExp.
     def enterParenExp(self, ctx:MathExprParser.ParenExpContext):
         pass

more_math/Parser/legacy/MathExprVisitor.py

@@ -124,6 +124,11 @@ class MathExprVisitor(ParseTreeVisitor):
         return self.visitChildren(ctx)
 
 
+    # Visit a parse tree produced by MathExprParser#LambdaExp.
+    def visitLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
+        return self.visitChildren(ctx)
+
+
     # Visit a parse tree produced by MathExprParser#TernaryExp.
     def visitTernaryExp(self, ctx:MathExprParser.TernaryExpContext):
         return self.visitChildren(ctx)
@@ -359,11 +364,6 @@ class MathExprVisitor(ParseTreeVisitor):
         return self.visitChildren(ctx)
 
 
-    # Visit a parse tree produced by MathExprParser#LambdaExp.
-    def visitLambdaExp(self, ctx:MathExprParser.LambdaExpContext):
-        return self.visitChildren(ctx)
-
-
     # Visit a parse tree produced by MathExprParser#ParenExp.
     def visitParenExp(self, ctx:MathExprParser.ParenExpContext):
         return self.visitChildren(ctx)

tests/test_unified_visitor_regressions.py

@@ -111,6 +111,24 @@ def test_stack_list_of_tensors_with_negative_dimension():
     assert torch.equal(result, torch.tensor([[1.0, 3.0], [2.0, 4.0]]))
 
 
+def test_lambda_binds_loosest():
+    assert _visit("f = (n) -> n <= 1 ? 1 : 2; f(0) + f(5)", {}) == 3.0
+    assert _visit("f = (g, n) -> n <= 1 ? 1 : n * g(g, n - 1); f(f, 4)", {}) == 24.0
+    assert _visit("h(p) -> p(3); h((x) -> x * 2)", {}) == 6.0
+    assert _visit("c = 1; f = c ? (x) -> x + 1 : (x) -> x - 1; f(5)", {}) == 6.0
+    assert _visit("f = (n) -> (n * 2); f(3)", {}) == 6.0
+    assert _visit("a = 2; (a + 3) * (a)", {}) == 10.0
+
+
+def test_script_has_to_end_with_a_value_not_a_lambda():
+    with pytest.raises(ValueError, match="ends with a lambda"):
+        _visit("(x) -> x * 2", {})
+    with pytest.raises(ValueError, match="ends with a lambda"):
+        _visit("f = (x) -> x * 2; f", {})
+    assert _visit("f = (x) -> x * 2; f(V0)", {"V0": 3.0}) == 6.0
+    assert _visit("V0 + V1", {"V0": 1.0, "V1": 2.0}) == 3.0
+
+
 def test_stack_separate_tensors_followed_by_dimension():
     first = torch.tensor([1.0, 2.0])
     second = torch.tensor([3.0, 4.0])