From 6ea548f03485a682d2d832f2bc5fe01b41329a23 Mon Sep 17 00:00:00 2001
From: Ozzie Gooen <ozzieagooen@gmail.com>
Date: Fri, 4 Oct 2024 17:02:37 -0700
Subject: [PATCH] Warning for if without else

---
 .../ai/__tests__/squiggleCodeWarnings_test.ts | 71 ++++++++++++++++++-
 .../ai/src/squiggle/squiggleCodeWarnings.ts   | 19 +++--
 2 files changed, 78 insertions(+), 12 deletions(-)

diff --git a/packages/ai/__tests__/squiggleCodeWarnings_test.ts b/packages/ai/__tests__/squiggleCodeWarnings_test.ts
index 611a021e12..fbaef5cac9 100644
--- a/packages/ai/__tests__/squiggleCodeWarnings_test.ts
+++ b/packages/ai/__tests__/squiggleCodeWarnings_test.ts
@@ -52,7 +52,7 @@ simulateExtinctionRisk() = {
     P_Misalignment = P_Misalignment_0 * exp(-K_align * E_align(t))
   })
 }
-    `;
+      `;
       const warnings = checkCapitalizedVariableNames(code);
       expect(warnings.length).toBe(5);
       expect(warnings[0].message).toContain(
@@ -149,10 +149,23 @@ Costs are provided per **full battery cycle** and per **hour of use**.",
       expect(warnings.length).toBe(0);
     });
   });
+
   describe("checkIfWithoutElse", () => {
     it("should detect if statements without else", () => {
       const code = `
-        if x > 5 then y = 10
+        z = if x > 5 then 10
+      `;
+      const warnings = checkIfWithoutElse(code);
+      expect(warnings.length).toBe(1);
+      expect(warnings[0].message).toContain(
+        "'if' statement without a corresponding 'else'"
+      );
+    });
+
+    it("should detect if statements without else, on next line", () => {
+      const code = `if year < 2023 then
+    throw("Year must be 2023 or later")
+
       `;
       const warnings = checkIfWithoutElse(code);
       expect(warnings.length).toBe(1);
@@ -160,6 +173,28 @@ Costs are provided per **full battery cycle** and per **hour of use**.",
         "'if' statement without a corresponding 'else'"
       );
     });
+
+    it("should not detect if statements with else", () => {
+      const code = `
+        x = if x > 5 then 10 else 20
+      `;
+      const warnings = checkIfWithoutElse(code);
+      expect(warnings.length).toBe(0);
+    });
+
+    it("should handle multiple if statements correctly", () => {
+      const code = `a = 3
+b = 5
+aa = if a > b then 5 else 10
+ab = if a > b then a
+      `;
+      const warnings = checkIfWithoutElse(code);
+      expect(warnings.length).toBe(1);
+      expect(warnings[0].message).toContain(
+        "'if' statement without a corresponding 'else'"
+      );
+      expect(warnings[0].lineNumber).toBe(4); // Assuming line numbers start at 1
+    });
   });
 
   describe("checkMultipleExpects", () => {
@@ -188,5 +223,37 @@ Costs are provided per **full battery cycle** and per **hour of use**.",
       expect(warnings.length).toBe(1);
       expect(warnings[0].message).toContain("Adjacent expect statements found");
     });
+
+    it("should not detect non-adjacent expect statements", () => {
+      const code = `
+        expect(x).toBe(5)
+
+        expect(y).toBe(10)
+      `;
+      const warnings = checkAdjacentExpectStatements(code);
+      expect(warnings.length).toBe(0);
+    });
+
+    it("should handle multiple adjacent expect statements correctly", () => {
+      const code = `
+        expect(a).toBe(1)
+        expect(b).toBe(2)
+        expect(c).toBe(3)
+      `;
+      const warnings = checkAdjacentExpectStatements(code);
+      expect(warnings.length).toBe(2);
+      expect(warnings[0].message).toContain("Adjacent expect statements found");
+      expect(warnings[1].message).toContain("Adjacent expect statements found");
+    });
+
+    it("should handle mixed expect and non-expect lines", () => {
+      const code = `
+        expect(a).toBe(1)
+        const b = 2
+        expect(c).toBe(3)
+      `;
+      const warnings = checkAdjacentExpectStatements(code);
+      expect(warnings.length).toBe(0);
+    });
   });
 });
diff --git a/packages/ai/src/squiggle/squiggleCodeWarnings.ts b/packages/ai/src/squiggle/squiggleCodeWarnings.ts
index 5f1c9f41b5..1a180714ba 100644
--- a/packages/ai/src/squiggle/squiggleCodeWarnings.ts
+++ b/packages/ai/src/squiggle/squiggleCodeWarnings.ts
@@ -278,25 +278,24 @@ export function checkCapitalizedVariableNames(code: string): Warning[] {
 export function checkIfWithoutElse(code: string): Warning[] {
   const warnings: Warning[] = [];
   const lines = code.split("\n");
-  const ifPattern = /^\s*if\s+.*\s+then\s+/;
-  const elsePattern = /^\s*else\s+/;
+  const ifPattern = /^\s*(\w+\s*=\s*)?if\s+/;
 
   for (let i = 0; i < lines.length; i++) {
     const trimmedLine = lines[i].trim();
     if (ifPattern.test(trimmedLine)) {
-      let hasElse = false;
-      // Check the next few lines for an else statement
-      for (let j = i + 1; j < Math.min(i + 5, lines.length); j++) {
-        if (elsePattern.test(lines[j].trim())) {
-          hasElse = true;
-          break;
-        }
+      let hasElse = trimmedLine.includes("else");
+
+      // Check the next two lines for 'else' if not found on the current line
+      for (let j = 1; j <= 2 && !hasElse && i + j < lines.length; j++) {
+        const nextLine = lines[i + j].trim();
+        hasElse = nextLine.includes("else");
       }
+
       if (!hasElse) {
         warnings.push({
           type: "CHECK_IF_WITHOUT_ELSE",
           lineNumber: i + 1,
-          message: `Line ${i + 1}: 'if' statement without a corresponding 'else'.`,
+          message: `Line ${i + 1}: 'if' statement without a corresponding 'else'. In Squiggle, every 'if' must have an 'else' clause.`,
         });
       }
     }