From 7be98b67cbad421a0b041d20e0f4620e70bd4cd6 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke_lau@icloud.com>
Date: Tue, 6 Aug 2019 15:27:58 +0100
Subject: [PATCH] Denormalize the type annotated ast, and tag stack values

---
 ast.scm       |  21 ++++----
 codegen.scm   |  55 +++++++++++++-------
 tests.scm     |  23 ++++++---
 typecheck.scm | 136 ++++++++++++++++++++++++++++++++++----------------
 4 files changed, 158 insertions(+), 77 deletions(-)

diff --git a/ast.scm b/ast.scm
index 4ca6eb1..6b546ac 100644
--- a/ast.scm
+++ b/ast.scm
@@ -17,6 +17,7 @@
       ('lambda 'lambda)      
       ('closure 'closure) ; only available in codegen
       ('static-string 'static-string) ; only available in codegen
+      ('stack 'stack) ; only available in codegen (tag that value is passed via stack)
       (else 'app)))
    ((builtin? x) 'builtin)
    ((symbol? x) 'var)
@@ -32,6 +33,7 @@
     ('app (map f x))
     ('lambda `(lambda ,(lambda-args x) ,(f (lambda-body x))))
     ('if `(if ,@(map f (cdr x))))
+    ('stack `(stack ,(cadr x) ,(map f (caddr x))))
     (else x)))
 
 (define (ast-collect f x)
@@ -46,7 +48,8 @@
 		     (inner (lambda-body x)))]
     ['if (append (f x)
 		 (flat-map inner (cdr x)))]
-    ['closure (flat-map inner (caddr x))]
+    ['stack (append (f x)
+		    (inner (caddr x)))]
     [else (f x)]))
 
 (define (ast-find p x)
@@ -69,6 +72,7 @@
     ['lambda (either (p x)
 		     (inner (lambda-body x)))]
     ['if (either (p x) (any inner (cdr x)))]
+    ['stack (either (p x) (inner (caddr x)))]
     [else (p x)]))
 
 (define (let-bindings e)
@@ -96,9 +100,13 @@
 	 (eqv? (car x) 'define)) 'define]
    [else 'expr]))
 
-(define (program-datas program)
-  (filter (lambda (x) (eqv? (statement-type x) 'data))
-	  program))
+
+					; (A ((foo (Int Bool))
+					;     (bar (Bool)))
+(define (program-data-layouts program)
+  (map (lambda (x) (cons (car x) (cdr x))) ; convert to assoc list
+       (map cdr (filter (lambda (x) (eqv? (statement-type x) 'data))
+			program))))
 
 (define (program-defines program)
   (filter (lambda (x) (eqv? (statement-type x) 'defines))
@@ -111,10 +119,6 @@
 	       program)))
 
 
-					; (A ((foo (Int Bool))
-					;     (bar (Bool)))
-
-(define data-layout cdr)
 
 					; gets both constructors and destructors
 					; (data A (foo Int Bool)
@@ -147,7 +151,6 @@
 	      [dtors (map (lambda (t i) (destructor ctor-name type-name t i))
 			  products
 			  (range 0 (length products)))])
-	 
 	 (cons maker (append dtors acc))))
      '()
      ctors)))
diff --git a/codegen.scm b/codegen.scm
index 2d5f4ac..55189cf 100644
--- a/codegen.scm
+++ b/codegen.scm
@@ -11,12 +11,12 @@
 
 (define wordsize 8)
 
-(define (type-size type env)
+(define (type-size data-layouts type)
 
   (define (adt-size adt)
     (let ([sizes
 	   (map (lambda (sum)
-		  (fold-left (lambda (acc x) (+ acc (type-size x)))
+		  (fold-left (lambda (acc x) (+ acc (type-size data-layouts x)))
 			     wordsize ; one word needed to store tag
 			     (cdr sum)))
 		(cdr adt))])
@@ -26,16 +26,21 @@
     ['Int wordsize]
     ['Bool wordsize]
     [else
-     (let ([adt (assoc type (env-adts env))])
+     (let ([adt (assoc type data-layouts)])
        (if adt
 	   (adt-size adt)
 	   (error #f "unknown size" type)))]))
 
+(define (on-stack? expr)
+  (case (ast-type expr)
+    ['stack (cadr expr)]
+    [else #f]))
+
 					; an environment consists of adt layouts in scope,
 					; and any bound variables.
 					; bound variables are an assoc list with their stack offset
 (define make-env list)
-(define env-adts car)
+(define env-data-layouts car)
 (define env-bindings cadr)
 
 (define (codegen-add xs si env)
@@ -160,7 +165,7 @@
 			      acc))
 		      (env-bindings env)
 		      comps)]
-		    [scc-env (make-env (env-adts env) scc-binding-offsets)])
+		    [scc-env (make-env (env-data-layouts env) scc-binding-offsets)])
 	       (for-each 
 		(lambda (name)
 		  (let ([expr (cadr (assoc name bindings))])
@@ -171,7 +176,7 @@
 			(codegen-expr expr
 				      inner-si
 				      (make-env
-				       (env-adts scc-env)
+				       (env-data-layouts scc-env)
 				       (cons (cons name 'self-captive)
 					     (env-bindings scc-env))))
 			(codegen-expr expr inner-si scc-env))
@@ -359,7 +364,7 @@
 
 (define (data-tor env e)
   (and (list? e)
-       (assoc (car e) (flat-map data-tors (env-adts env)))))
+       (assoc (car e) (flat-map data-tors (env-data-layouts env)))))
 
 (define (codegen-data-tor e si env)
 
@@ -400,7 +405,8 @@
 	    (codegen-call (car e) (cdr e) si env)))))
 
 					; this is a builtin being passed around as a variable
-    ('builtin (emit "movq $~a, %rax" (builtin-id e)))
+					; this should have been converted to a closure!
+    ('builtin (error #f "passing about a builtin!" e))
 
     ('let (codegen-let (let-bindings e)
 		       (let-body e)
@@ -419,7 +425,20 @@
 
     (else (error #f "don't know how to codegen this"))))
 
-
+					; takes in a expr annotated with types and returns a type-less AST
+					; with stack values wrapped
+(define (annotate-stack-values data-layout ann-e)
+  (define (struct-type? type)
+    (assoc type data-layout))
+  (define (strip e)
+    (ast-traverse strip (ann-expr e)))
+  (let* ([e (ann-expr ann-e)]
+	 [type (ann-type ann-e)])
+    (if (struct-type? type)
+	`(struct ,(type-size data-layout type) ,(ast-traverse strip e))
+	(ast-traverse (lambda (x)
+			(annotate-stack-values data-layout x))
+		      e))))
 
 (define (free-vars prog)
   (define bound '())
@@ -583,15 +602,17 @@
 (define (codegen program)
   (set! cur-label 0)
   (set! cur-lambda 0)
-  (let* ([body (program-body program)]
+  (let* ([data-layouts (program-data-layouts program)]
 
-	 [data-layouts (map data-layout (program-datas program))]
+	 [type-annotated (annotate-types program)]
+	 [stack-annotated (annotate-stack-values data-layouts
+						 type-annotated)]
 	 
-	 (extract-res-0 (extract-strings body))
-	 (strings (car extract-res-0))
-	 (extract-res-1 (extract-lambdas (cdr extract-res-0)))
-	 (lambdas (car extract-res-1))
-	 (xform-prog (cdr extract-res-1)))
+	 (strings-res (extract-strings stack-annotated))
+	 (strings (car strings-res))
+	 (lambdas-res (extract-lambdas (cdr strings-res)))
+	 (lambdas (car lambdas-res))
+	 (xform-prog (cdr lambdas-res)))
 
     (emit "\t.global _start")
     (emit "\t.text")
@@ -647,6 +668,6 @@
 ; %r8  = 5th arg
 ; %r9  = 6th arg
 
-; on darwin, the syscall is offset by 0x2000000
+; on darwin, unix/posix syscalls are offset by 0x2000000 (syscall classes)
 ; https://opensource.apple.com/source/xnu/xnu-2782.20.48/bsd/kern/syscalls.master
 ; documentation for most syscalls: /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk/usr/include/sys
diff --git a/tests.scm b/tests.scm
index 4e14fb6..87ed2a6 100644
--- a/tests.scm
+++ b/tests.scm
@@ -43,18 +43,18 @@
   (let ((str (read-file "/tmp/test-output.txt")))
     (test str output)))
 
-(test (data-tors (data-layout '(data A
-			 (foo Int Bool)
+(test (data-tors '(A . ((foo Int Bool)
 			(bar Bool))))
       '((foo . (constructor . (abs Int (abs Bool A))))
 	(foo~0 . (0 . (abs A Int)))
 	(foo~1 . (1 . (abs A Bool)))
 	(bar . (constructor . (abs Bool A)))
-	(bar~0 . (0 . (abs A Bool)))))
+	(bar~0 . (0 . (
+
+		       abs A Bool)))))
 
 (test (data-tors-env
-       (data-layout '(data A
-			    (foo Int Bool)
+       '(A . ((foo Int Bool)
 	      (bar Bool))))
       '((foo . (abs Int (abs Bool A)))
 	(foo~0 . (abs A Int))
@@ -154,9 +154,9 @@
 	   (- y x))))
 
       '((let ()
-	  ((let ((x 42 : Int)
-		 (y ((((+ : (abs Int (abs Int Int))) (1 : Int)) : (abs Int Int)) (x : Int)) : Int))
-	     (((((- : (abs Int (abs Int Int))) (y : Int)) : (abs Int Int)) (x : Int)) : Int))))))
+	  ((let ((x (42 : Int))
+		 (y (((+ : (abs Int (abs Int Int))) (1 : Int) (x : Int)) : Int)))
+	     (((- : (abs Int (abs Int Int))) (y : Int) (x : Int)) : Int)) : Int)) : Int))
 
 (test-expr '(+ 1 2) 3)
 (test-expr '(bool->int (= 2 0)) 0)
@@ -224,3 +224,10 @@
 ;; 				(print "b")
 ;; 				 (f (- m 1))))])
 ;; 			 (f 10))) "ababababab")
+
+					; adts and pattern matching
+
+(test-prog '((data (A [foo Int]))
+	     (let ([(foo x) (foo 42)])
+	       x))
+	   42)
diff --git a/typecheck.scm b/typecheck.scm
index cfb30cc..e398606 100644
--- a/typecheck.scm
+++ b/typecheck.scm
@@ -131,7 +131,7 @@
 	       scc-env)]
 
 	 [annotated-bindings (append (cdr acc) ; the previous annotated bindings
-				     (map cons
+				     (map list
 					  comps
 					  (map caddr type-results)))])
       (cons new-env annotated-bindings)))
@@ -146,9 +146,47 @@
 	 [let-type (cadr (last body-results))]
 	 [cs (fold-left (lambda (acc cs) (constraint-merge acc cs)) '() (map car body-results))]
 
-	 [annotated `((let ,annotated-bindings ,@(map caddr body-results)))])
+	 [annotated `((let ,annotated-bindings ,@(map caddr body-results)) : ,let-type)])
     (list cs let-type annotated)))
 
+(define (check-app env x)
+  (if (eqv? (car x) (cadr x))
+					; recursive function (f f)
+					; TODO: what about ((f a) f)????
+      (let* ([func-type (env-lookup env (car x))]
+	     [return-type (fresh-tvar)]
+	     [other-func-type `(abs ,func-type ,return-type)]
+	     [cs (~ func-type other-func-type)]
+	     [resolved-return-type (substitute cs return-type)]
+
+	     [annotated `(((,(car x) : ,func-type)
+			   (,(cadr x) : ,func-type)) : ,resolved-return-type)])
+	(list cs resolved-return-type annotated)))
+
+					; regular function
+  (let* ([arg-type-res (check env (cadr x))]
+	 [arg-type (cadr arg-type-res)]
+	 [func-type-res (check env (car x))]
+	 [func-type (cadr func-type-res)]
+	 
+					; f ~ a -> t0
+	 [func-c (~
+		  (substitute (car arg-type-res) func-type)
+		  `(abs ,arg-type ,(fresh-tvar)))]
+	 [cs (constraint-merge
+	      (constraint-merge func-c (car arg-type-res))
+	      (car func-type-res))]
+	 
+	 [resolved-func-type (substitute cs func-type)]
+	 [resolved-return-type (caddr resolved-func-type)]
+
+	 [annotated `((,(caddr func-type-res)
+		       ,(caddr arg-type-res)) : ,resolved-return-type)])
+
+    (if (abs? resolved-func-type)
+	(let ((return-type (substitute cs (caddr resolved-func-type))))
+	  (list cs return-type annotated))
+	(error #f "not a function"))))
 
 ; returns a list (constraints type annotated)
 (define (check env x)
@@ -199,51 +237,14 @@
 
 		     [lambda-type `(abs ,resolved-arg-type ,(cadr body-type-res))]
 
-		     ; TODO: do we need to annotate the lambda argument?
-		     [annotated `(lambda (,(lambda-arg x)) ,(caddr body-type-res))])
+		     [annotated `((lambda (,(lambda-arg x)) ,(caddr body-type-res)) : ,lambda-type)])
 		
 		(list (car body-type-res) ; constraints
 		      lambda-type  ; type
 		      annotated)))
 
 	  
-	  ('app ; (f a)
-	   (if (eqv? (car x) (cadr x))
-					; recursive function (f f)
-	       (let* ([func-type (env-lookup env (car x))]
-		      [return-type (fresh-tvar)]
-		      [other-func-type `(abs ,func-type ,return-type)]
-		      [cs (~ func-type other-func-type)]
-		      [resolved-return-type (substitute cs return-type)]
-
-		      [annotated `(((,(car x) : ,func-type)
-				    (,(cadr x) : ,func-type)) : ,resolved-return-type)])
-		 (list cs resolved-return-type annotated)))
-
-					; regular function
-	   (let* ([arg-type-res (check env (cadr x))]
-		  [arg-type (cadr arg-type-res)]
-		  [func-type-res (check env (car x))]
-		  [func-type (cadr func-type-res)]
-		  
-					; f ~ a -> t0
-		  [func-c (~
-			   (substitute (car arg-type-res) func-type)
-			   `(abs ,arg-type ,(fresh-tvar)))]
-		  [cs (constraint-merge
-		       (constraint-merge func-c (car arg-type-res))
-		       (car func-type-res))]
-		  
-		  [resolved-func-type (substitute cs func-type)]
-		  [resolved-return-type (caddr resolved-func-type)]
-
-		  [annotated `((,(caddr func-type-res)
-				,(caddr arg-type-res)) : ,resolved-return-type)])
-
-	     (if (abs? resolved-func-type)
-		 (let ((return-type (substitute cs (caddr resolved-func-type))))
-		   (list cs return-type annotated))
-		 (error #f "not a function")))))))
+	  ('app (check-app env x)))))
     ;; (display "result of ")
     ;; (display x)
     ;; (display ":\n\t")
@@ -254,14 +255,63 @@
     res))
 
 (define (init-adts-env prog)
-  (flat-map data-tors-env (map data-layout (program-datas prog))))
+  (flat-map data-tors-env (program-data-layouts prog)))
 
 					; we typecheck the lambda calculus only (only single arg lambdas)
 (define (typecheck prog)
   (cadr (check (init-adts-env prog) (normalize (program-body prog)))))
 
+
+					; before passing annotated types onto codegen
+					; we need to restore the pre-normalization structure
+					; (this is important for function arity etc)
+(define (denormalize orig normed)
+
+  (define (collapse-lambdas n x)
+    (case n
+      [0 x]
+      [else
+       (let* ([inner-lambda (lambda-body (ann-expr x))]
+	      [arg (lambda-arg (ann-expr x))]
+	      [inner-collapsed (ann-expr (collapse-lambdas (- n 1) inner-lambda))])
+	 `((lambda ,(cons arg (lambda-args inner-collapsed))
+	     ,(lambda-body inner-collapsed)) : ,(ann-type x)))]))
+
+  (define (collapse-apps n x)
+    (case n
+      [-1 (error #f "nullary functions not handled yet")]
+      [0 x]
+      [else
+       (let* ([inner-app (car (ann-expr x))]
+	      [inner-collapsed (collapse-apps (- n 1) inner-app)])
+	 `(,(append (ann-expr inner-collapsed) (cdr (ann-expr x))) : ,(ann-type x)))]))
+
+  (case (ast-type orig)
+    ['lambda
+	(let ([collapsed (collapse-lambdas (- (length (lambda-args orig)) 1) normed)])
+	  `((lambda ,(lambda-args (ann-expr collapsed))
+	      ,(denormalize (lambda-body orig)
+			    (lambda-body (ann-expr collapsed)))) : ,(ann-type collapsed)))]
+    ['app
+     (let ([collapsed (collapse-apps (- (length orig) 2) normed)])
+       `(,(map (lambda (o n) (denormalize o n)) orig (ann-expr collapsed))
+	 : ,(ann-type collapsed)))]
+    ['let
+	`((let ,(map (lambda (o n) (list (car o) (denormalize (cadr o) (cadr n))))
+		     (let-bindings orig)
+		     (let-bindings (ann-expr normed)))
+	    ,@(map (lambda (o n) (denormalize o n))
+		   (let-body orig)
+		   (let-body (ann-expr normed)))) : ,(ann-type normed))]
+    ['if `((if ,@(map denormalize (cdr orig) (cdr (ann-expr normed))))
+	   : (ann-type normed))]
+    [else normed]))
+
+(define ann-expr car)
+(define ann-type caddr)
 (define (annotate-types prog)
-  (caddr (check (init-adts-env prog) (normalize (program-body prog)))))
+  (denormalize (program-body prog)
+  (caddr (check (init-adts-env prog) (normalize (program-body prog))))))
 
   
 					; returns a list of constraints
-- 
2.30.2