#|
exec sbcl --core /usr/share/sbcl/sbcl-iterate --noinform --disable-debugger --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)
(shadow 'defmacro)

(defstruct env globals constants builtins funcs vars labels macros)
(defparameter *env* (make-env))

(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 intern-macro (macro-name params expander-fn env)
  (push (list macro-name params expander-fn) (env-macros env)))
(cl:defmacro defmacro (mname (&rest syntax) &body body)
  "sugar defmacro"
  `(intern-macro ',mname ',syntax
    (lambda (form env)
      (declare (ignorable env))
      (destructuring-bind ,syntax (cdr form) ,@body))
    *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 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-constant (var typ val env)
  (push (list var typ val) (env-constants env)))
(defun bind-label (label env)
  (push label (env-labels env)))
(defun tagp (tag form)
  (and (consp form) (eq tag (car form))))

(declaim (ftype (function (t t) list) compile))
(declaim (ftype (function (list t) list) compile-sequence))
(defun var-get (var env)
  (cond
    ((eq 't var) '(i32.const 1))
    ((assoc var (env-constants env))
      (destructuring-bind (typ val) (cdr (assoc var (env-constants env)))
        `(,(symcat typ :.const) ,val)))
    ((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 (tagp 'locals 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-defconstant (expr env)
  (destructuring-bind (var typ val) expr
    (bind-constant var typ val env)
    "")) ; empty string is how we emit no code
(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.u64* :i64.store) (set.u32* :i32.store) (set.u16* :i32.store16) (set.u8* :i32.store8)
              (set.i64* :i64.store) (set.i32* :i32.store) (set.i16* :i32.store16) (set.i8* :i32.store8))))
      `(,store ,(compile ptr env) ,(compile val env)))))
(defun special-load (expr env)
  (destructuring-bind (getmode ptr) expr
    (let* ((load
            (ecase getmode
              (get* :i32.load)
              (get.u64* :i64.load) (get.u32* :i32.load) (get.u16* :i32.load16_u) (get.u8* :i32.load8_u)
              (get.i64* :i64.load) (get.i32* :i32.load) (get.i16* :i32.load16_s) (get.i8* :i32.load8_s))))
      `(,load ,(compile ptr env)))))
(defun special-coerce (expr env)
  (destructuring-bind (val dst-type src-type) expr
    (declare (ignore dst-type src-type)) ;; todo
    `(:i32.wrap_i64 ,(compile val env))
    ))
(defun special-i64 (expr env)
  "(i64 <64-bit number>) - produce (i64.const n)"
  (declare (ignore env))
  `(i64.const ,(car expr)))
(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 special-vprogn (expr env)
  (destructuring-bind (typ &body body) expr
    `(block (result ,typ) ,@(compile-sequence body env))))
(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 macro-while (form env)
  (declare (ignore env))
  (destructuring-bind (condition &body body) (cdr form)
    (let* ((break (genlabel :break)) (head (genlabel :head)))
      `(block ,break (loop ,head
         (br-if ,break (zerop ,condition))
         ,@body
         (br ,head))))))
(defun macro-forever (form env)
  (declare (ignore env))
  (destructuring-bind (&body body) (cdr form)
    (let* ((break (genlabel :break)) (head (genlabel :head)))
      `(block ,break (loop ,head
         ,@body
         (br ,head))))))
(defun macro-break (form env)
  (destructuring-bind (&optional (label (cadr (env-labels env)))) (cdr form)
    `(br ,label)))
(defun macro-break-if (form env)
  (destructuring-bind (condition &optional (label (cadr (env-labels env)))) (cdr form)
    `(br-if ,label ,condition)))
(defun macro-continue (form env)
  (destructuring-bind (&optional (label (car (env-labels env)))) (cdr form)
    `(br ,label)))
(defun macro-and-rec (form typ)
  (if (cdr form)
    `(if (result ,typ) ,(car form) ,(macro-and-rec (cdr form) typ) 0)
    (car form)))
(defun macro-or-rec (form typ)
  (if (cdr form)
    `(vprogn ,typ (set scortmp ,(car form)) (if (result ,typ) scortmp scortmp ,(macro-or-rec (cdr form) typ)))
    (car form)))
(defun macro-and (form env)
  (declare (ignore env))
  (if (null (cdr form))
    't
    (macro-and-rec (cdr form) 'i32)))
(defun macro-or (form env)
  (declare (ignore env))
  (macro-or-rec (cdr form) 'i32))
(defun macro-aref (form env)
  "(aref ptr idx [multiplier = 4])"
  (declare (ignore env))
  (destructuring-bind (ptr idx &optional (multiplier 4)) (cdr form)
    (let* ((desc (cdr (assoc multiplier '((1 get.u8* 0)(2 get.u16* 1)(4 get* 2)(8 get.i64* 3))))))
      (unless desc (error "sugar: I don't have a getter for objects of size ~d. Use (+ ptr (* idx multiplier)) yourself." multiplier))
      `(,(car desc) (+ ,ptr (<< ,idx ,(cadr desc)))))))
(defun macro-aset (form env)
  "(aset ptr idx value [multiplier = 4]) <--- later to be inferred with setf"
  (declare (ignore env))
  (destructuring-bind (ptr idx val &optional (multiplier 4)) (cdr form)
    (let* ((desc (cdr (assoc multiplier '((1 set.u8* 0)(2 set.u16* 1)(4 set* 2)(8 set.i64* 3))))))
      (unless desc (error "sugar: I don't have a setter for objects of size ~d. Use (+ ptr (* idx multiplier)) yourself." multiplier))
      `(,(car desc) (+ ,ptr (<< ,idx ,(cadr desc))) ,val))))
(defmacro when (condition &body body)
  `(if ,condition (progn ,@body)))
(defmacro not (expr) `(zerop ,expr))
(defmacro unless (condition &body body)
  `(if (not ,condition) (progn ,@body)))

(defun compile-atom (expr env)
  (declare (ignore env))
  (typecase expr
    (number `(i32.const ,expr))
    (null '(i32.const 0))
    (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))
    (defconstant (special-defconstant (cdr expr) env))
    (if (special-if (cdr expr) env))
    (progn (special-progn (cdr expr) env))
    (vprogn (special-vprogn (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))
    (coerce (special-coerce (cdr expr) env))
    (i64 (special-i64 (cdr expr) env))
    ((set* set.i32* set.u32* set.i64* set.u64* set.i16* set.u16* set.i8* set.u8*) (special-store expr env))
    ((get* get.i32* get.u32* get.i64* get.u64* get.i16* get.u16* get.i8* 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-case* :downcase) (*print-pretty* t) (*print-miser-width* 120))
        (format t "~a~%" (compile form env)))))
  (format t "  (global $scortmp (mut i32) (i32.const 0))~%")
  (format t ")~%")
  (format *error-output* "summary: globals:~d constants:~d functions:~d macros:~d~%"
    (length (env-globals env))
    (length (env-constants 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.mul nil env)
  (intern-builtin '/ :i32.div_u nil env)
  (intern-builtin '% :i32.rem_u 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 'zerop :i32.eqz nil env)
  (intern-builtin '= :i32.eq nil env)
  (intern-builtin '/= :i32.ne nil env)
  (intern-builtin '< :i32.lt_u nil env)
  (intern-builtin '<= :i32.le_u nil env)
  (intern-builtin '> :i32.gt_u 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)
  (intern-macro 'while nil #'macro-while env)
  (intern-macro 'forever nil #'macro-forever env)
  (intern-macro 'break nil #'macro-break env)
  (intern-macro 'break-if nil #'macro-break-if env)
  (intern-macro 'continue nil #'macro-continue env)
  (intern-macro 'and nil #'macro-and env)
  (intern-macro 'or nil #'macro-or env)
  (intern-macro 'aref nil #'macro-aref env)
  (intern-macro 'aset nil #'macro-aset env)

  (bind-global 'scortmp 'i32 env)
  )

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