Validate returns inside if statements

common/addons/structure-terrascript-loader/src/main/java/com/dfsek/terra/addons/terrascript/semanticanalysis/TypeChecker.java

@@ -2,7 +2,6 @@ package com.dfsek.terra.addons.terrascript.semanticanalysis;
 
 import java.util.List;
 import java.util.Optional;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.dfsek.terra.addons.terrascript.Environment;
 import com.dfsek.terra.addons.terrascript.Environment.Symbol;
@@ -167,19 +166,9 @@ public class TypeChecker implements Visitor<TypedExpr>, Stmt.Visitor<TypedStmt>
     
     @Override
     public TypedStmt visitFunctionDeclarationStmt(Stmt.FunctionDeclaration stmt) {
-        AtomicBoolean hasReturn = new AtomicBoolean(false);
+        TypedStmt.Block body = new TypedStmt.Block(stmt.body.statements.stream().map(s -> s.accept(this)).toList());
-        TypedStmt.Block body = new TypedStmt.Block(stmt.body.statements.stream().map(s -> {
+        boolean hasReturn = alwaysReturns(body, stmt);
-            TypedStmt bodyStmt = s.accept(this);
+        if(stmt.returnType != Type.VOID && !hasReturn) {
-            if(bodyStmt instanceof TypedStmt.Return ret) {
-                hasReturn.set(true);
-                if(ret.value.type != stmt.returnType)
-                    errorHandler.add(new InvalidTypeException(
-                            "Return statement must match function's return type. Function '" + stmt.identifier + "' expects " +
-                            stmt.returnType + ", found " + ret.value.type + " instead", s.position));
-            }
-            return bodyStmt;
-        }).toList());
-        if(stmt.returnType != Type.VOID && !hasReturn.get()) {
             errorHandler.add(
                     new InvalidFunctionDeclarationException("Function body for '" + stmt.identifier + "' does not contain return statement",
                                                             stmt.position));
@@ -187,6 +176,23 @@ public class TypeChecker implements Visitor<TypedExpr>, Stmt.Visitor<TypedStmt>
         return new TypedStmt.FunctionDeclaration(stmt.identifier, stmt.parameters, stmt.returnType, body, scopedIdentifier(stmt.identifier, stmt.getSymbol()));
     }
     
+    private boolean alwaysReturns(TypedStmt stmt, Stmt.FunctionDeclaration function) {
+        if(stmt instanceof TypedStmt.Return ret) {
+            if(ret.value.type != function.returnType)
+                errorHandler.add(new InvalidTypeException(
+                        "Return statement must match function's return type. Function '" + function.identifier + "' expects " +
+                        function.returnType + ", found " + ret.value.type + " instead", function.position));
+            return true;
+        } else if (stmt instanceof TypedStmt.If ifStmt) {
+            return alwaysReturns(ifStmt.trueBody, function) &&
+                   ifStmt.elseIfClauses.stream().map(Pair::getRight).allMatch(s -> alwaysReturns(s, function)) &&
+                   ifStmt.elseBody.map(body -> alwaysReturns(body, function)).orElse(false); // If else body is not defined then statement does not always return
+        } else if (stmt instanceof TypedStmt.Block block) {
+            return block.statements.stream().anyMatch(s -> alwaysReturns(s, function));
+        }
+        return false;
+    }
+    
     @Override
     public TypedStmt visitVariableDeclarationStmt(Stmt.VariableDeclaration stmt) {
         TypedExpr value = stmt.value.accept(this);

common/addons/structure-terrascript-loader/src/test/java/semanticanalysis/SemanticAnalyzerTest.java

@@ -92,9 +92,51 @@ public class SemanticAnalyzerTest {
         testValid("fun returnVoid() { return (); }");
     }
     
-    //Not implemented yet @Test
+    @Test
-    public void testControlFlowAnalysis() {
+    public void testFunctionReturnControlFlowAnalysis() {
+        // Non-void returning function bodies must contain at least one statement that always returns
+        testInvalid("""
+                    fun returnsNum(): num {
+                    }
+                    """, InvalidFunctionDeclarationException.class);
         
+        testValid("""
+                  fun returnsNum(): num {
+                      return 1;
+                  }
+                  """);
+        
+        testValid("""
+                  fun returnsNum(): num {
+                      1 + 1;
+                      return 1;
+                  }
+                  """);
+        
+        // Statements after the first always-return-statement are unreachable, unreachable code is legal
+        testValid("""
+                  fun returnsNum(): num {
+                      return 1;
+                      1 + 1; // Unreachable
+                  }
+                  """);
+        
+        // Void returning functions can omit returns
+        testValid("""
+                  fun returnsNothing() {
+                  }
+                  """);
+        
+        // Returns can still be explicitly used for void returning functions
+        testValid("""
+                    fun returnsNum(p: bool) {
+                        if (p) {
+                            return;
+                        }
+                    }
+                    """);
+        
+        // If all if-statement bodies always return, then the statement is considered as always returning
         testValid("""
                   fun returnsNum(p: bool): num {
                       if (p) {
@@ -105,31 +147,84 @@ public class SemanticAnalyzerTest {
                   }
                   """);
         
+        testValid("""
+                  fun returnsNum(p1: bool, p2: bool): num {
+                      if (p1) {
+                          return 1;
+                      } else if (p2) {
+                          return 2;
+                      } else {
+                          return 3;
+                      }
+                  }
+                  """);
         
-        // All paths of execution must return a value if return type is not void
+        // If no else body is defined, an if-statement does not always return, therefore the function does not contain any always-return-statements
         testInvalid("""
                     fun returnsNum(p: bool): num {
                         if (p) {
                             return 1;
                         }
                     }
-                    """, null);
+                    """, InvalidFunctionDeclarationException.class);
         
-        // Not all paths require explicit return for void functions, implicitly returns void at end of execution
+        testInvalid("""
+                    fun returnsNum(p1: bool, p2: bool): num {
+                        if (p1) {
+                            return 1;
+                        } else if (p2) {
+                            return 2;
+                        }
+                    }
+                    """, InvalidFunctionDeclarationException.class);
+        
+        // Nested ifs should work
         testValid("""
-                    fun returnsNum(p: bool) {
+                  fun returnsNum(p1: bool, p2: bool): num {
-                        if (p) {
+                      if (p1) {
-                            return;
+                          if (p2) {
+                              return 1;
+                          } else {
+                              return 2;
+                          }
+                      } else {
+                          return 3;
                       }
                   }
                   """);
         
-        // Static analysis of if statement always being true
+        testInvalid("""
-        testValid("""
+                  fun returnsNum(p1: bool, p2: bool): num {
-                  fun returnsNum(): num {
+                      if (p1) {
-                      if (true) {
+                          if (p2) {
                               return 1;
                           }
+                          // No else clause here, so will not always return
+                      } else {
+                          return 3;
+                      }
+                  }
+                  """, InvalidFunctionDeclarationException.class);
+        
+        // If-statement may not always return but a return statement after it means function will always return
+        testValid("""
+                    fun returnsNum(p: bool): num {
+                        if (p) {
+                            return 1;
+                        }
+                        return 2;
+                    }
+                    """);
+        
+        // Same applies when statements are swapped
+        testValid("""
+                    fun returnsNum(p: bool): num {
+                        return 1;
+                        // Unreachable
+                        if (p) {
+                            return 2;
+                        }
+                    }
                     """);
     }