Fix implicit return from if/else and return inside try/or

if/else blocks as the last expression in a function or lambda now
correctly return the branch value instead of nil:

  def classify(x)
    if x > 10
      "big"
    else
      "small"
    end
  end
  classify(20)  # => "big" (was nil)

return statements inside try/or error handlers no longer cause
a compile error. They now set the handler result value:

  matches = try find_stuff(input) or err
    return nil  # sets matches = nil (previously failed to compile)
  end

Author: rubiojr

compiler/codegen.go

@@ -52,6 +52,7 @@ type codeGen struct {
 	currentFunc     *ast.FuncDef         // current function being generated (for type lookups)
 	varTypeScope    string               // override scope key for varType lookups (test/bench blocks)
 	inSpawn         int                  // nesting depth of spawn blocks (>0 means inside spawn)
+	inTryHandler    int                  // nesting depth of try/or handler defers (>0 means inside handler)
 	lambdaDepth     int                  // nesting depth of lambda bodies (>0 means inside fn)
 	lambdaScopeBase []int                // scope index at each lambda entry (stack)
 	lambdaOuterFunc []*ast.FuncDef       // enclosing function at each lambda entry (stack)
@@ -978,16 +979,14 @@ func (g *codeGen) writeFunc(f *ast.FuncDef) error {
 	g.inFunc = true
 	hasImplicitReturn := false
 	for i, s := range f.Body {
-		// Implicit return: last expression in function body becomes the return value.
+		// Implicit return: last expression or if/else in function body becomes the return value.
 		if i == len(f.Body)-1 {
-			if es, ok := s.(*ast.ExprStmt); ok {
-				g.emitLineDirective(es.SourceLine)
-				expr, err := g.exprString(es.Expression)
-				if err != nil {
-					return err
-				}
-				g.writef("return %s\n", expr)
-				hasImplicitReturn = true
+			handled, allCovered, err := g.writeLastStmtAs(s, "return %s\n")
+			if err != nil {
+				return err
+			}
+			if handled {
+				hasImplicitReturn = allCovered
 				continue
 			}
 		}
@@ -1220,6 +1219,51 @@ func (g *codeGen) writeExprStmt(e *ast.ExprStmt) error {
 func (g *codeGen) writeIf(i *ast.IfStmt) error {
 	// Pre-declare variables assigned in any branch so they're visible
 	// after the if block (Ruby-like scoping: if/else doesn't create a new scope).
+	g.predeclareIfVars(i)
+
+	cond, err := g.exprString(i.Condition)
+	if err != nil {
+		return err
+	}
+	g.writef("if %s {\n", g.condExpr(cond, i.Condition))
+	g.indent++
+	for _, s := range i.Body {
+		if err := g.writeStmt(s); err != nil {
+			return err
+		}
+	}
+	g.indent--
+	for _, ec := range i.ElsifClauses {
+		cond, err := g.exprString(ec.Condition)
+		if err != nil {
+			return err
+		}
+		g.writef("} else if %s {\n", g.condExpr(cond, ec.Condition))
+		g.indent++
+		for _, s := range ec.Body {
+			if err := g.writeStmt(s); err != nil {
+				return err
+			}
+		}
+		g.indent--
+	}
+	if len(i.ElseBody) > 0 {
+		g.writeln("} else {")
+		g.indent++
+		for _, s := range i.ElseBody {
+			if err := g.writeStmt(s); err != nil {
+				return err
+			}
+		}
+		g.indent--
+	}
+	g.writeln("}")
+	return nil
+}
+
+// predeclareIfVars pre-declares variables assigned in any branch of an if/else
+// so they're visible after the if block (Ruby-like scoping).
+func (g *codeGen) predeclareIfVars(i *ast.IfStmt) {
 	var allBranches []ast.Statement
 	allBranches = append(allBranches, i.Body...)
 	for _, ec := range i.ElsifClauses {
@@ -1237,44 +1281,91 @@ func (g *codeGen) writeIf(i *ast.IfStmt) error {
 			g.declareVar(name)
 		}
 	}
+}
+
+// writeLastStmtAs tries to write a statement as an implicit return or assignment.
+// format is the format string for the last expression (e.g. "return %s\n" or "r = %s\n").
+// Returns handled=true if the statement was written, allCovered=true if all code paths
+// produce a value (e.g. if/elsif/else with all branches handled).
+func (g *codeGen) writeLastStmtAs(s ast.Statement, format string) (handled bool, allCovered bool, err error) {
+	switch st := s.(type) {
+	case *ast.ExprStmt:
+		g.emitLineDirective(st.StmtLine())
+		expr, err := g.exprString(st.Expression)
+		if err != nil {
+			return false, false, err
+		}
+		g.writef(format, expr)
+		return true, true, nil
+	case *ast.IfStmt:
+		allCovered, err := g.writeIfWithLastAction(st, format)
+		if err != nil {
+			return false, false, err
+		}
+		return true, allCovered, nil
+	default:
+		return false, false, nil
+	}
+}
+
+// writeIfWithLastAction writes an if/elsif/else block where the last expression
+// in each branch is formatted with the given format string (for implicit returns).
+func (g *codeGen) writeIfWithLastAction(i *ast.IfStmt, format string) (bool, error) {
+	g.predeclareIfVars(i)
 
 	cond, err := g.exprString(i.Condition)
 	if err != nil {
-		return err
+		return false, err
 	}
 	g.writef("if %s {\n", g.condExpr(cond, i.Condition))
 	g.indent++
-	for _, s := range i.Body {
-		if err := g.writeStmt(s); err != nil {
-			return err
-		}
+	if err := g.writeBodyWithLastAction(i.Body, format); err != nil {
+		return false, err
 	}
 	g.indent--
 	for _, ec := range i.ElsifClauses {
 		cond, err := g.exprString(ec.Condition)
 		if err != nil {
-			return err
+			return false, err
 		}
 		g.writef("} else if %s {\n", g.condExpr(cond, ec.Condition))
 		g.indent++
-		for _, s := range ec.Body {
-			if err := g.writeStmt(s); err != nil {
-				return err
-			}
+		if err := g.writeBodyWithLastAction(ec.Body, format); err != nil {
+			return false, err
 		}
 		g.indent--
 	}
+	allCovered := false
 	if len(i.ElseBody) > 0 {
 		g.writeln("} else {")
 		g.indent++
-		for _, s := range i.ElseBody {
-			if err := g.writeStmt(s); err != nil {
-				return err
-			}
+		if err := g.writeBodyWithLastAction(i.ElseBody, format); err != nil {
+			return false, err
 		}
 		g.indent--
+		allCovered = true
 	}
 	g.writeln("}")
+	return allCovered, nil
+}
+
+// writeBodyWithLastAction writes a list of statements, converting the last
+// expression into the given format (e.g. "return %s\n").
+func (g *codeGen) writeBodyWithLastAction(body []ast.Statement, format string) error {
+	for i, s := range body {
+		if i == len(body)-1 {
+			handled, _, err := g.writeLastStmtAs(s, format)
+			if err != nil {
+				return err
+			}
+			if handled {
+				continue
+			}
+		}
+		if err := g.writeStmt(s); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
@@ -1476,6 +1567,18 @@ func (g *codeGen) rangeIntExpr(e ast.Expr) string {
 }
 
 func (g *codeGen) writeReturn(r *ast.ReturnStmt) error {
+	// Inside a try/or handler defer, return sets the handler result (r).
+	if g.inTryHandler > 0 {
+		if r.Value != nil {
+			expr, err := g.exprString(r.Value)
+			if err != nil {
+				return err
+			}
+			g.writef("r = %s\n", expr)
+		}
+		g.writeln("return")
+		return nil
+	}
 	// Inside a spawn block, return EXPR must assign to t.result and
 	// use a bare return (the goroutine closure has no return value).
 	if g.inSpawn > 0 {
@@ -2145,27 +2248,30 @@ func (g *codeGen) tryExpr(e *ast.TryExpr) (string, error) {
 	handlerCode, cerr := g.captureOutput(func() error {
 		g.pushScope()
 		g.declareVar(e.ErrVar)
+		g.inTryHandler++
 
 		for i, s := range e.Handler {
 			isLast := i == len(e.Handler)-1
 			if isLast {
-				// Last statement: if it's a bare expression, assign to r (return value)
-				if es, ok := s.(*ast.ExprStmt); ok {
-					val, verr := g.exprString(es.Expression)
-					if verr != nil {
-						g.popScope()
-						return verr
-					}
-					g.writef("r = %s\n", val)
+				// Last statement: if it's a bare expression or if/else, assign to r (return value)
+				handled, _, herr := g.writeLastStmtAs(s, "r = %s\n")
+				if herr != nil {
+					g.inTryHandler--
+					g.popScope()
+					return herr
+				}
+				if handled {
 					continue
 				}
 			}
 			if werr := g.writeStmt(s); werr != nil {
+				g.inTryHandler--
 				g.popScope()
 				return werr
 			}
 		}
 
+		g.inTryHandler--
 		g.popScope()
 		return nil
 	})
@@ -2274,15 +2380,13 @@ func (g *codeGen) fnExpr(e *ast.FnExpr) (string, error) {
 		for i, s := range e.Body {
 			isLast := i == len(e.Body)-1
 			if isLast {
-				// Last statement: if it's a bare expression, make it the return value
-				if es, ok := s.(*ast.ExprStmt); ok {
-					g.emitLineDirective(es.StmtLine())
-					val, verr := g.exprString(es.Expression)
-					if verr != nil {
-						restoreLambda()
-						return verr
-					}
-					g.writef("return %s\n", val)
+				// Last statement: if it's a bare expression or if/else, make it the return value
+				handled, _, herr := g.writeLastStmtAs(s, "return %s\n")
+				if herr != nil {
+					restoreLambda()
+					return herr
+				}
+				if handled {
 					continue
 				}
 			}

rats/core/09_try_or_test.rugo

@@ -31,3 +31,9 @@ rats "shell fallback with try"
   test.assert_eq(result["status"], 0)
   test.assert_contains(result["output"], "continued")
 end
+
+rats "return inside try/or handler sets handler result"
+  result = test.run("rugo run rats/core/fixtures/bug_try_or_return_type.rugo")
+  test.assert_eq(result["status"], 0)
+  test.assert_eq(result["output"], "1")
+end

rats/core/54_implicit_return_test.rugo

@@ -81,3 +81,37 @@ rats "explicit return still works"
   test.assert_eq(result["status"], 0)
   test.assert_eq(result["output"], "small\nbig")
 end
+
+rats "implicit return from if/else in function"
+  source = <<~RUGO
+    def classify(x)
+      if x > 10
+        "big"
+      else
+        "small"
+      end
+    end
+    puts(classify(20))
+    puts(classify(5))
+  RUGO
+  result = eval.run(source)
+  test.assert_eq(result["status"], 0)
+  test.assert_eq(result["output"], "big\nsmall")
+end
+
+rats "implicit return from if/else in lambda"
+  source = <<~RUGO
+    arr = [1, 2, 3, nil]
+    result = arr.map(fn(x)
+      if x != nil
+        x * 2
+      else
+        nil
+      end
+    end)
+    puts result
+  RUGO
+  result = eval.run(source)
+  test.assert_eq(result["status"], 0)
+  test.assert_eq(result["output"], "[2 4 6 <nil>]")
+end

rats/core/fixtures/bug_try_or_return_type.rugo

@@ -0,0 +1,17 @@
+def find_stuff(query)
+  return [1, 2, 3]
+end
+
+def resolve(input)
+  matches = try find_stuff(input) or err
+    return nil
+  end
+
+  if len(matches) == 0
+    return nil
+  end
+
+  return matches[0]
+end
+
+puts resolve("test")