From a457cd3bb5ce9366db3ca0731a07abc50ecbc1f3 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke_lau@icloud.com>
Date: Thu, 15 Aug 2019 17:21:13 +0100
Subject: [PATCH] Fix bindings in ADT pattern matching codegen

---
 codegen.scm   | 38 +++++++++++++++++++++++++++++---------
 tests.scm     |  7 +++++++
 typecheck.scm |  2 +-
 3 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/codegen.scm b/codegen.scm
index a70ef6c..6203cfd 100644
--- a/codegen.scm
+++ b/codegen.scm
@@ -69,6 +69,9 @@
 (define make-env list)
 (define env-data-layouts car)
 (define env-bindings cadr)
+(define (env-append-bindings env bindings)
+  (make-env (env-data-layouts env)
+	    (append bindings (env-bindings env))))
 
 (define (codegen-add xs si env)
   (define (go ys)
@@ -444,18 +447,32 @@
 					; checks if equal and returns assoc list of bindings
   (define (check-equal jne-label type inner-offset x)
 
+					; TODO: tidy this up! comparibles and binds could be merged
 					; (foo a 2 (bar x)) -> ((2 Int 1) ((bar x) A 2))
 					; sum: foo
     (define (comparibles sum)
-      (if (not (list? sum))
-	  '()
       (let ([product-types (cdr (assoc sum (cdr (assoc type dls))))])
+	(if (null? product-types)
+	    '()
 	    (filter (lambda (x) (not (eqv? 'var (ast-type (car x)))))
 		    (map (lambda (x t i) (list x t i))
 			 (cdr x)
 			 product-types
 			 (range 0 (length product-types)))))))
     
+    (define (binds sum)
+      (let ([product-types (cdr (assoc sum (cdr (assoc type dls))))])
+	(if (null? product-types)
+	    '()
+	    (filter (lambda (x) (eqv? 'var (ast-type (car x))))
+		    (map (lambda (x i)
+			   (cons x
+				 (- inner-offset
+				    (data-product-offset dls type sum i))))
+			 (cdr x)
+			 (range 0 (length (cdr x))))))))
+	     
+
     (let ([sums (assoc type dls)])
       (if sums
 	  (let* ([sum (if (list? x) (car x) x)] ; can sometimes be a singleton
@@ -463,13 +480,15 @@
 					; the tag is at the top (beginning) of the adt on the stack
 	    (emit "cmpq $~a, ~a(%rbp)" tag inner-offset)
 	    (emit "jne ~a" jne-label)
+
+	    (append (binds sum)
 	     (flat-map
-	     (lambda (cmpx cmpt cmpi)
+	     (lambda (cmp) ; cmp = (x type index)
 	       (check-equal jne-label
-			    cmpt
-			    (- inner-offset (data-product-offset dls type sum i))
-			    cmpx))
-	     (comparibles sum)))
+			    (cadr cmp)
+			    (- inner-offset (data-product-offset dls type sum (caddr cmp)))
+			    (car cmp)))
+	     (comparibles sum))))
 	  (if (eqv? 'var (ast-type x))
 	      (list (cons x inner-offset))
 	      (begin
@@ -482,8 +501,9 @@
 	   [expr (cadr case)]
 	   [next-section-label (fresh-label)]
 	   [inner-si (- si (type-size dls type))]
-	   [new-env (append (check-equal next-section-label type si match)
-			    env)])
+	   [new-env (env-append-bindings env
+		     (check-equal next-section-label type si match))])
+
 
       (codegen-expr expr inner-si new-env)
       (emit "jmp ~a" exit-label)
diff --git a/tests.scm b/tests.scm
index 0c11e89..a3f4913 100644
--- a/tests.scm
+++ b/tests.scm
@@ -300,3 +300,10 @@
 		 [c 1])))
 	   2)
 
+(test-prog '((data Foo [foo Int Int] [bar Bool])
+	     (case (foo 42 12)
+	       [(foo 20 x) 0]
+	       [(foo 42 x) x]
+	       [(bar x) 0]))
+	   12)
+	       
diff --git a/typecheck.scm b/typecheck.scm
index 064f65f..c6b90e7 100644
--- a/typecheck.scm
+++ b/typecheck.scm
@@ -219,7 +219,7 @@
 
 	 [resolved-type (substitute case-expr-equality-cs (car case-expr-types))]
 
-	 [annotated `((case (,(case-switch x) : ,switch-type)
+	 [annotated `((case ,(caddr switch-type-res)
 			,@(map (lambda (c e et)
 				 `(,c ((,e : ,et))))
 			       (map car (case-cases x))
-- 
2.30.2