#|
exec sbcl --noinform --core /usr/share/sbcl/sbcl-iterate --load $0 --end-toplevel-options "$@"
|#
;; --- sugar compiler - compiles sugar to wasm/WAT

(defpackage :sugar
  (:use :cl :iterate)) ;; oh the humanity! why didn't he use ASDF?
(in-package :sugar)

(setf *print-pretty* nil)

(shadow 'compile)
(shadow 'load)
(shadow 'functionp)

(defstruct env globals builtins funcs vars labels macros)

(defun symcat (&rest syms)
  (intern (format nil "~:@(~{~a~}~)" syms)))
(defun symstr (&rest syms)
  (format nil "\"~{~a~}\"" syms))
(defparameter next-label 0)
(defun genlabel (prefix)
  (format nil "~a~d" prefix (incf next-label)))

(defun macrop (expr env)
  (assoc (car expr) (env-macros env)))
(defun expando (expr env)
  (let* ((macro-desc (assoc (car expr) (env-macros env))))
    (assert macro-desc)
    (funcall (nth 2 macro-desc) expr env)))

(defun export-function (fname env)
  (declare (ignore env))
  (unless (eq 'export (car fname)) (error "sugar: expected (export <funcname>) but found ~s" fname))
  `(export ,(format nil "\"~a\"" (cadr fname)))
  )
(defun functionp (fname env)
  (assoc fname (env-funcs env)))
(defun builtinp (fname env)
  (assoc fname (env-builtins env)))
(defun intern-function (key-name params env)
  (push (list key-name params) (env-funcs env)))
(defun intern-builtin (key-name wat-name params env)
  (push (list key-name wat-name params) (env-builtins env)))
(defun intern-macro (macro-name params expander-fn env)
  (push (list macro-name params expander-fn) (env-macros env)))
(defun bind-var (var typ env)
  (push (list var typ) (env-vars env)))
(defun bind-global (var typ env)
  (push (list var typ) (env-globals env)))
(defun bind-label (label env)
  (push label (env-labels env)))

(declaim (ftype (function (t t) list) compile))
(declaim (ftype (function (list t) list) compile-sequence))
(defun var-get (var env)
  (cond
    ((assoc var (env-globals env)) `(global.get ,(symcat :$ var)))
    ((assoc var (env-vars env)) `(local.get ,(symcat :$ var)))
    (t (error "sugar: undefined variable ~s" var))))
(defun var-set (var val env)
  (cond
    ((assoc var (env-globals env)) `(global.set ,(symcat :$ var) ,(compile val env)))
    ((assoc var (env-vars env)) `(local.set ,(symcat :$ var) ,(compile val env)))
    (t (error "sugar: undefined variable ~s" var))))

(defun special-defun-sig (expr env)
  (destructuring-bind (fname (&rest params) &body body) expr
    (let* ((wat-params (iter (for (var typ) on params by #'cddr) (collect `(param ,(symcat :$ var) ,typ))))
           (rtype
              (if (eq :result (car body))
                (prog1
                  `((result ,(cadr body)))
                  (setf body (cddr body)))))
           (key-name fname)
           (wat-fname (symcat :$ fname)))
      (when (consp fname)
        (setf key-name (cadr fname))
        (setf wat-fname (export-function fname env)))
      (intern-function key-name params env)
      (values
        `(func ,wat-fname ,@wat-params ,@rtype)
        params
        body))))
(defun defun-locals (clauses env)
  (let* (vardefs)
    (append
      (iter (for (typ . vars) in clauses)
        (appending
          (iter (for var in vars)
            (cond
              ((consp var)
                (push var vardefs)
                (bind-var (car var) typ env)
                (collect `(local ,(symcat :$ (car var)) ,typ)))
              (t
                (bind-var var typ env)
                (collect `(local ,(symcat :$ var) ,typ)))))))
      (iter (for (var val) in vardefs) (collect `(local.set ,(symcat :$ var) ,(compile val env)))))))
(defun special-defun (expr env)
  (multiple-value-bind (head params body) (special-defun-sig expr env)
    (let* ((entry-scope (env-vars env)) (maybe-locals (car body)) locals wat)
      (iter (for (var typ) on params by #'cddr) (bind-var var typ env))
      (when (and (consp maybe-locals) (eq 'locals (car maybe-locals)))
        (setf body (cdr body))
        (setf locals (defun-locals (cdr maybe-locals) env)))
      (setf wat `(,@head ,@locals ,@(compile-sequence body env)))
      (setf (env-vars env) entry-scope)
      wat)))
(defun special-import (expr env)
  "(import \"host\" \"print\" (func print (i i32)))"
  (destructuring-bind (module func sig) expr
    `(import ,(symstr module) ,(symstr func) ,(special-defun-sig (cdr sig) env))
    ))
(defun special-memory (expr env)
  (declare (ignore env))
  "(memory (export \"mem\") 1)
   (memory (import \"js\" \"mem\") 10)"
  (ecase (caar expr)
    (export `(memory (export ,(symstr (nth 1 (car expr)))) ,(cadr expr)))
    (import `(memory (import ,(symstr (nth 1 (car expr))) ,(symstr (nth 2 (car expr)))) ,(cadr expr)))
    ))
(defun special-global (expr env)
  (destructuring-bind (var typ val) expr
    (prog1
      `(global ,(symcat :$ var) ,typ ,(compile val env))
      (bind-global var typ env))))
(defun special-set (expr env)
  (destructuring-bind (var val) expr
    (var-set var val env)))
(defun special-store (expr env)
  (destructuring-bind (setmode ptr val) expr
    (let* ((store (ecase setmode (set* :i32.store) (set.i16* :i32.store16) (set.i8* :i32.store8))))
      `(,store ,(compile ptr env) ,(compile val env)))))
(defun special-load (expr env)
  (destructuring-bind (setmode ptr) expr
    (let* ((load (ecase setmode (get* :i32.load) (get.u8* :i32.load8_u))))
      `(,load ,(compile ptr env)))))
(defun special-block (expr env)
  (destructuring-bind (block label &body body) expr
    (let* ((entry-scope (env-labels env)) wat)
      (bind-label label env)
      (setf wat `(,block ,(symcat :$ label) ,@(compile-sequence body env)))
      (setf (env-labels env) entry-scope)
      wat
      )))
(defun special-progn (expr env)
  (destructuring-bind (&body body) expr
    `(block ,@(compile-sequence body env))))
(defun tagp (tag form)
  (and (consp form) (eq tag (car form))))
(defun special-if (expr env)
  (let* (maybe-result)
    (when (tagp 'result (car expr))
      (setf maybe-result (list (car expr)))
      (setf expr (cdr expr)))
    (destructuring-bind (condition then &optional else) expr
      `(if ,@maybe-result ,(compile condition env) ,(compile then env) ,@(if else (list (compile else env)))))))
(defun special-br (expr env)
  (destructuring-bind (label) expr
    (unless (member label (env-labels env)) (error "sugar: undefined branch label: ~a" label))
    `(br ,(symcat :$ label))))
(defun special-br-if (expr env)
  (destructuring-bind (label condition) expr
    (unless (member label (env-labels env)) (error "sugar: undefined branch label: ~a" label))
    `(br_if ,(symcat :$ label) ,(compile condition env))))

(defun macro-inc (form env)
  (declare (ignore env))
  (destructuring-bind (var &optional (delta 1)) (cdr form)
    `(set ,var (+ ,var ,delta))))
(defun macro-for (form env)
  (declare (ignore env))
  (destructuring-bind ((idx start end &optional (step 1)) &body body) (cdr form)
    (let* ((break (genlabel :break)) (head (genlabel :head)))
      `(block ,break
         (set ,idx ,start)
         (loop ,head
           (br-if ,break (= ,end ,idx))
           ,@body
           (inc ,idx ,step)
           (br ,head))))))

(defun compile-atom (expr env)
  (declare (ignore env))
  (typecase expr
    (number `(i32.const ,expr))
    (otherwise `(ATOM-TODO ,expr))
    ))
(defun compile-special (expr env)
  (case (car expr)
    (defun (special-defun (cdr expr) env))
    (import (special-import (cdr expr) env))
    (memory (special-memory (cdr expr) env))
    (global (special-global (cdr expr) env))
    (if (special-if (cdr expr) env))
    (progn (special-progn (cdr expr) env))
    (block (special-block expr env))
    (loop (special-block expr env))
    (br (special-br (cdr expr) env))
    (br-if (special-br-if (cdr expr) env))
    (set (special-set (cdr expr) env))
    ((set* set.i16* set.i8*) (special-store expr env))
    ((get* get.u8*) (special-load expr env))
    ))
(defun compile-sequence (exprs env)
  (iter (for x in exprs) (collect (compile x env))))
(defun compile-funcall (expr env)
  (cond
    ((builtinp (car expr) env) `(,(cadr (builtinp (car expr) env)) ,@(iter (for x in (cdr expr)) (collect (compile x env)))))
    ((functionp (car expr) env) `(call ,(symcat :$ (car expr)) ,@(iter (for x in (cdr expr)) (collect (compile x env)))))
    (t (error "sugar: undefined function ~s called" (car expr)))))

(defun compile (expr env)
  (typecase expr
      ((or null number string character)
        (compile-atom expr env))
      (symbol (var-get expr env))
      (cons
        (if (macrop expr env)
          (compile (expando expr env) env)
          (or (compile-special expr env) (compile-funcall expr env))))))

(defun load (fname env)
  (format t "(module~%")
  (with-open-file (f fname)
    (iter
      (for form = (read f nil nil))
      (while form)
      (let ((*print-pretty* t) (*print-miser-width* 120))
        (format t "~a~%" (compile form env)))))
  (format t ")~%")
  (format *error-output* "summary: globals:~d functions:~d macros:~d~%"
    (length (env-globals env))
    (length (env-funcs env))
    (length (env-macros env))))

(defun compiler-init (env)
  (intern-builtin 'memory.size :memory.size nil env)
  (intern-builtin 'memory.grow :memory.grow nil env)

  (intern-builtin '+ :i32.add nil env)
  (intern-builtin '- :i32.sub nil env)
  (intern-builtin '& :i32.and nil env)
  (intern-builtin 'bor :i32.or nil env)
  (intern-builtin '^ :i32.xor nil env)
  (intern-builtin '>> :i32.shr_u nil env)
  (intern-builtin '<< :i32.shl nil env)

  (intern-builtin '= :i32.eq nil env)
  (intern-builtin 'zerop :i32.eqz nil env)
  (intern-builtin '/= :i32.ne nil env)
  (intern-builtin '>= :i32.ge_u nil env)

  (intern-builtin 'return :return nil env)
  (intern-builtin 'select :select nil env)
  (intern-builtin 'drop :drop nil env)

  (intern-macro 'inc nil #'macro-inc env)
  (intern-macro 'for nil #'macro-for env)
  )

(let* ((env (make-env)))
  (compiler-init env)
  (load (cadr sb-ext:*posix-argv*) env))
(sb-ext:exit :code 0)