From 061f7cd9efa96f5d4e7206ec89931f9fd8421a6c Mon Sep 17 00:00:00 2001
From: Luke Lau <luke_lau@icloud.com>
Date: Thu, 15 Aug 2019 16:51:36 +0100
Subject: [PATCH] Add case statement codegen and singleton ADTs

---
 codegen.scm | 130 ++++++++++++++++++++++++++++++++++++++++++++++------
 tests.scm   |   6 +--
 2 files changed, 119 insertions(+), 17 deletions(-)

diff --git a/codegen.scm b/codegen.scm
index dd79507..a70ef6c 100644
--- a/codegen.scm
+++ b/codegen.scm
@@ -36,9 +36,9 @@
 	   (error #f "unknown size" type)))]))
 
 					; returns the size of an expression's result in bytes
-(define (expr-size e)
+(define (expr-size dls e)
   (if (eqv? (ast-type e) 'stack)
-      (cadr e)
+      (type-size dls (cadr e))
       wordsize))
 
 (define (on-stack? expr)
@@ -47,11 +47,20 @@
     [else #f]))
 
 					; does a movsq for something on the stack
+					; src points to the start stack index, but not the top of that index
+					; likewise for dst
+					; | ...  |
+					; +------+ <-- to here
+					; | tag0 |
+					; +------+ <-- src (size = 16)
+					; |  42  |
+					; +------+ <-- start copying from here...
 (define (emit-stack-copy src dst size)
-  (emit "leaq ~a(%rbp), %rsi" (- src size))
-  (emit "leaq ~a(%rbp), %rdi" (- dst size))
+  (let ([size-to-copy (- size wordsize)])
+    (emit "leaq ~a(%rbp), %rsi" (- src size-to-copy))
+    (emit "leaq ~a(%rbp), %rdi" (- dst size-to-copy))
     (emit "movq $~a, %rcx" (/ size wordsize))
-  (emit "rep movsq"))	  
+    (emit "rep movsq")))
   
 
 					; an environment consists of adt layouts in scope,
@@ -164,8 +173,12 @@
 
   ;;   (fold-left emit-binding (cons env '()) scc))))
 					; assoc map of binding name to size
+
+
+  (define dls (env-data-layouts env))
+  
   (define stack-sizes
-    (map (lambda (binding) (cons (car binding) (expr-size (cadr binding))))
+    (map (lambda (binding) (cons (car binding) (expr-size dls (cadr binding))))
 	 bindings))
 
 					; assoc map of binding name to offset
@@ -205,7 +218,7 @@
 	       (for-each 
 		(lambda (name)
 		  (let* ([expr (cadr (assoc name bindings))]
-			 [size (expr-size expr)])
+			 [size (expr-size dls expr)])
 		    (emit "## generating ~a with scc-env ~a" name scc-env)
 		    (if (self-captive-closure? name expr)
 					; if self-captive, insert a flag into the environment to let
@@ -232,15 +245,30 @@
 	      body)))
 
 (define (codegen-var e si env)
-  (let* ([stack-size (on-stack? e)]
-	 [name (if (on-stack? e) (caddr e) e)]
-	 [stack-offset (cdr (assoc name (env-bindings env)))])
-    (when (not stack-offset)
+  (let* ([stack-type (on-stack? e)]	 
+	 [name (if stack-type (caddr e) e)]
+
+	 [dls (env-data-layouts env)]
+	 
+	 [singleton? (and stack-type
+			(assoc name
+			       (cdr (assoc stack-type dls))))]
+				    
+	 [stack-offset (assoc name (env-bindings env))])
+    (when (and (not stack-offset) (not singleton?))
       (error #f (format "Variable ~a is not bound" name)))
 
     (if (on-stack? e)
-	(emit-stack-copy stack-offset si stack-size)
-	(emit "movq ~a(%rbp), %rax" stack-offset))))
+	(if singleton?
+					; singletons don't need to be in the environment
+					; just copy over the tag
+	    (emit "movq $~a, ~a(%rbp)"
+		  (data-sum-tag (env-data-layouts env)
+				stack-type
+				name)
+		  si)
+	    (emit-stack-copy (cdr stack-offset) si (type-size dls stack-type)))
+	(emit "movq ~a(%rbp), %rax" (cdr stack-offset)))))
 
 (define cur-lambda 0)
 (define (fresh-lambda)
@@ -408,6 +436,79 @@
     (codegen-expr else si env)
     (emit "~a:" exit-label)))
 
+(define (codegen-case switch cases si env)
+  (define dls (env-data-layouts env))
+  (define exit-label (fresh-label))
+
+
+					; checks if equal and returns assoc list of bindings
+  (define (check-equal jne-label type inner-offset x)
+
+					; (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))))])
+	    (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)))))))
+
+    (let ([sums (assoc type dls)])
+      (if sums
+	  (let* ([sum (if (list? x) (car x) x)] ; can sometimes be a singleton
+		 [tag (data-sum-tag dls type sum)])
+					; 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)
+	    (flat-map
+	     (lambda (cmpx cmpt cmpi)
+	       (check-equal jne-label
+			    cmpt
+			    (- inner-offset (data-product-offset dls type sum i))
+			    cmpx))
+	     (comparibles sum)))
+	  (if (eqv? 'var (ast-type x))
+	      (list (cons x inner-offset))
+	      (begin
+		(emit "cmp $~a, ~a(%rbp)" x inner-offset)
+		(emit "jne ~a" jne-label)
+		'() )))))
+  
+  (define (codegen-adt-match type case)
+    (let* ([match (car case)]
+	   [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)])
+      
+      (codegen-expr expr inner-si new-env)
+      (emit "jmp ~a" exit-label)
+
+      (emit "~a:" next-section-label)))
+
+  (define (codegen-literal-match case)
+    (let ([next-section-label (fresh-label)])
+      (emit "cmpq $~a, %rax" (car case))
+      (emit "jne ~a" next-section-label)
+      (codegen-expr (cadr case) si env)
+      (emit "jmp ~a" exit-label)
+      (emit "~a:" next-section-label)))
+  
+					; generate the switch
+					; (and store it on the stack if not a stack value)
+  (codegen-expr switch si env)
+
+  (if (eqv? 'stack (ast-type switch))
+					; adt pattern match
+      (for-each (lambda (x) (codegen-adt-match (cadr switch) x)) cases)
+      (for-each codegen-literal-match cases))
+  (emit "~a:" exit-label))
+      
+
 (define (data-tor env e)
   (if (not (list? e)) #f    
       (assoc (car e) (flat-map data-tors (env-data-layouts env)))))
@@ -531,6 +632,7 @@
     ('var (codegen-var e si env))
 
     ('if (codegen-if (cadr e) (caddr e) (cadddr e) si env))
+    ('case (codegen-case (case-switch e) (case-cases e) si env))
     
     ('bool-literal (emit "movq $~a, %rax" (if e 1 0)))
     ('int-literal (emit "movq $~a, %rax" e))
@@ -555,7 +657,7 @@
   (let* ([e (ann-expr ann-e)]
 	 [type (ann-type ann-e)])
     (if (stack-type? type)
-	`(stack ,(type-size data-layouts type)
+	`(stack ,type
 		,(ast-traverse (lambda (x) (annotate-stack-values data-layouts x)) e))
 	(ast-traverse (lambda (x)
 			(annotate-stack-values data-layouts x))
diff --git a/tests.scm b/tests.scm
index 7185032..0c11e89 100644
--- a/tests.scm
+++ b/tests.scm
@@ -295,8 +295,8 @@
 (test-prog '((data Foo [a] [b] [c])
 	     (let ([x b])
 	       (case x
-		 [a b]
-		 [b a]
-		 [c x])))
+		 [a 3]
+		 [b 2]
+		 [c 1])))
 	   2)
 
-- 
2.30.2