From 4db3f8e9c8cadfeeda9941473a322e3110c7387d Mon Sep 17 00:00:00 2001
From: Tim Nelson <tbnelson@gmail.com>
Date: Wed, 13 Mar 2024 15:44:21 -0400
Subject: [PATCH] add: better error for ifte mismatch of args

---
 forge/lang/expander.rkt              |  2 +-
 forge/sigs-structs.rkt               | 50 +++++++++++++++++-----------
 forge/tests/error/loc/ifte-error.frg | 13 ++++++++
 forge/tests/error/main.rkt           |  1 +
 4 files changed, 46 insertions(+), 20 deletions(-)
 create mode 100644 forge/tests/error/loc/ifte-error.frg

diff --git a/forge/lang/expander.rkt b/forge/lang/expander.rkt
index a745bfae..7b4f638c 100644
--- a/forge/lang/expander.rkt
+++ b/forge/lang/expander.rkt
@@ -1126,7 +1126,7 @@
          (first xs)]         
         [else
          (raise-forge-error
-          #:msg (format "Ill-formed block: expected either one expression or any number of formulas")
+          #:msg (format "Ill-formed block: expected either one expression or any number of formulas; got ~a" xs)
           #:context stx)]))
 
 ; Block : /LEFT-CURLY-TOK Expr* /RIGHT-CURLY-TOK
diff --git a/forge/sigs-structs.rkt b/forge/sigs-structs.rkt
index a6d34bd7..df431976 100644
--- a/forge/sigs-structs.rkt
+++ b/forge/sigs-structs.rkt
@@ -652,11 +652,25 @@ Returns whether the given run resulted in sat or unsat, respectively.
 ; or expression form (which has its own AST node). Avoid exponential
 ; blowup from chained IFTEs by expanding to a chain of function calls.
 (define (ifte-disambiguator info a b c)
-  (if (node/formula? b)
-      (&&/info info
-               (=>/info info a b)
-               (=>/info info (!/info info a) c))
-      (ite/info info a b c)))
+  (unless (node/formula? a)
+    (raise-forge-error
+     #:msg ("If-then-else needed a boolean-valued formula for its first argument; got ~a." (pretty-type-of a))
+     #:context a))
+  (cond
+    ; It's a formula if-then-else
+    [(and (node/formula? b) (node/formula? c))
+     (&&/info info
+              (=>/info info a b)
+              (=>/info info (!/info info a) c))]
+    ; It's an expression if-then-else
+    [(and (node/expr? b) (node/expr? c))
+     (ite/info info a b c)]
+    ; It's an error
+    [else
+     (raise-forge-error #:msg (format "If-then-else needed consistent types (either both formulas or both expressions) for its true and false branches, but got (~a) and (~a)."
+                                      (pretty-type-of b) (pretty-type-of c))
+                        #:context info)]))
+
 (define-syntax (ifte stx)
   (syntax-parse stx 
     [(_ (~optional (#:lang check-lang) #:defaults ([check-lang #''checklangNoCheck])) a b c) (quasisyntax/loc stx
@@ -722,19 +736,17 @@ Returns whether the given run resulted in sat or unsat, respectively.
             (->/info info univ b)))
 
 (define (domain-check<: a b loc) 
-  (let ([src-line (source-location-line loc)]
-        [src-col (source-location-column loc)]
-        [src-span (source-location-span loc)])
-    (unless (equal? (node/expr-arity b)
-                    (@+ 1 (node/expr-arity a))) 
-                    (raise-user-error (format "<: argument has incorrect arity (~a vs. ~a) in ~a <: ~a on line ~a, column ~a, span ~a" 
-                    (node/expr-arity a) (node/expr-arity b) (deparse a) (deparse b) src-line src-col src-span)))))
+  (unless (equal? (node/expr-arity b)
+                  (@+ 1 (node/expr-arity a))) 
+    (raise-forge-error
+     #:msg (format "<: argument has incorrect arity (~a vs. ~a) in ~a <: ~a" 
+                   (node/expr-arity a) (node/expr-arity b) (deparse a) (deparse b))
+     #:context loc)))
 
 (define (domain-check:> a b loc) 
-  (let ([src-line (source-location-line loc)]
-        [src-col (source-location-column loc)]
-        [src-span (source-location-span loc)])
-    (unless (equal? (node/expr-arity a)
-                    (@+ 1 (node/expr-arity b))) 
-                    (raise-user-error (format ":> argument has incorrect arity (~a vs. ~a) in ~a :> ~a on line ~a, column ~a, span ~a" 
-                    (node/expr-arity a) (node/expr-arity b) (deparse a) (deparse b) src-line src-col src-span)))))
+  (unless (equal? (node/expr-arity a)
+                  (@+ 1 (node/expr-arity b))) 
+    (raise-forge-error
+     #:msg (format ":> argument has incorrect arity (~a vs. ~a) in ~a :> ~a" 
+                   (node/expr-arity a) (node/expr-arity b) (deparse a) (deparse b))
+     #:context loc)))
diff --git a/forge/tests/error/loc/ifte-error.frg b/forge/tests/error/loc/ifte-error.frg
new file mode 100644
index 00000000..bf770e63
--- /dev/null
+++ b/forge/tests/error/loc/ifte-error.frg
@@ -0,0 +1,13 @@
+#lang forge/temporal 
+
+-- Test for srcloc preservation in if-then-else case
+
+sig Node {edges: set Node}
+one sig A, B, C extends Node {}
+
+pred if_example {
+     {some A.edges.edges} 
+       =>   {some B.edges.edges} 
+       else {C.edges.edges}
+} 
+test expect { {if_example} is sat}
\ No newline at end of file
diff --git a/forge/tests/error/main.rkt b/forge/tests/error/main.rkt
index 8b740dfc..0c054bfe 100644
--- a/forge/tests/error/main.rkt
+++ b/forge/tests/error/main.rkt
@@ -39,6 +39,7 @@
     ;;;;;;; Source locations ;;;;;;;
     (list "./loc/sig_use_loc_error.frg" #rx"sig_use_loc_error.frg:7:39") ; vs. reachable
     (list "./loc/field_use_loc_error.frg" #rx"field_use_loc_error.frg:7:29")   ; vs. reachable
+    (list "./loc/ifte-error.frg" #rx"ifte-error.frg:9:5.*If-then-else needed consistent types")  
    
     (list "piecewise-bind-repeat.frg" #rx"rebinding detected")
     (list "piecewise-bind-combine.frg" #rx"may not be combined with complete bounds")