; ssc: the sufficiently simple compiler

; The goal of ssc is to allow simple prefix expressions to be compiled into 65816 code that
; would run at least as fast or faster than the equivalent threaded Forth code. Complex
; optimizations are a non-goal; if you want to tune the generated code, go ahead and write
; the assembly directly.

; * All values are 16-bit integers, like Forth or BCPL.
; * Data and return addresses are mixed on one stack, unlike Forth.
; * Function calls take a fixed number of arguments, and return 0 or 1 results. The compiler enforces arity checking.
; * To call a function taking arguments [arg1 arg2 arg3], all 3 arguments should be pushed to the stack before calling.
;   When the function takes control, the stack should look like this:
;     arg1 arg2 arg3 return-address
; * The caller is responsible for removing all arguments from the stack once the function returns.
; * The caller is responsible for preserving the A, X and Y registers, if this is desirable.
; * If the function returns a value, it is stored in the A register.
; * If a function returns no result, it is not obliged to preserve the A register.
; * Multitasking is achieved by overlapping the D and S registers on the same 256-byte page of memory.
;   Yielding to a new task involves saving the S register, setting the D register to the new task's page,
;   then setting the S register to the saved value in the old task.
; * Useful task-local "registers" are kept at the beginning of the page, and the stack grows down from the end of the page.
;   * DP register list:
;     * bank-local address (16 bits)
;     * long address (24 bits)
;     * Last suspended value of S
;     * Mailbox
;     * Pointer to next task

; Compiler notes:
; Expressions are of the form [:function arg1 arg2 arg3]
; args are either strings (symbols) or numbers

; TODO:
; * implement read / write (pointers)
; * implement write (locals)
; * implement loops
; * implement "getters" (subroutine that runs when referenced by name without an explicit call)

(local Object (require :core.object))
(local lume (require :lib.lume))
(local Ssc (Object:extend))
(local Prg (require :asm.asm))

(fn Ssc.new [self ?opts]
  (local opts (or ?opts {}))
  (set self.prg (or opts.prg (Prg.new nil :65816)))
  (set self.forms (lume.clone (or opts.forms self.__index.forms)))
  (set self.functions {})
  (set self.locals []))

(fn Ssc.push [self ?local]
  (table.insert self.locals (or ?local [:tmp]))
  [:pha])
(fn Ssc.drop [self ?local]
  (match (. self.locals (length self.locals))
    [:tmp] (when (not= ?local nil) (error (.. "Internal stack error: expected " ?local ", got temporary")))
    loc    (when (not= ?local loc) (error (.. "Internal stack error: expected " (or ?local "temporary") " temporary, got " loc))))
  (tset self.locals (length self.locals) nil)
  [:ply])

(fn countiter [minmax ?max]
  (let [min (if ?max minmax 1)
        max (or ?max minmax)]
    (fn [_ iprev]
      (let [i (if iprev (+ iprev 1) min)]
        (when (<= i max) i)))))

(fn pairoff [l]
  (fn [_ iprev] (let [i (if iprev (+ iprev 2) 1)]
                          (when (< i (length l)) (values i (. l i) (. l (+ i 1)))))))

(fn boolop [self left right branch]
  (let [compiled-left    (self:compile-expr left)
        push-left        (when (not= (type right) :number) (self:push))
        compiled-compare (if (not push-left) [:cmp right]
                           [:block push-left (self:compile-expr right) [:cmp 0 :s]])
        drop-left (when push-left (self:drop))]
    [:block
      compiled-left
      compiled-compare
      [branch :-true-]
      [:lda 0]
      [:bra :-finished-]
    :-true-
      [:lda 0xffff]
    :-finished-
      drop-left]))

(set Ssc.forms
  {:asm (fn [self ...] [:block ...])
   :org (lambda [self org] (set self.org (self.prg:org org)))
   :do (fn [self ...] (lume.concat [:block] (icollect [i (countiter (select :# ...))] (self:compile-expr (select i ...)))))
   :let (fn [self bindings ...]
          (lume.concat [:block]
                       (icollect [_ symbol expr (pairoff bindings)]
                          [:block (self:compile-expr expr) (self:push symbol)])
                       [(self:compile-expr [:do ...])]
                       (icollect [i-half (countiter (/ (length bindings) 2))]
                          (self:drop (. bindings (- (length bindings) (* i-half 2) -1))))))
   :+ (lambda [self first ...]
        (local block [:block (self:compile-expr first)])
        (for [i 1 (select :# ...)]
          (let [val (select i ...)]
            (table.insert block
              (match val
                 1 [:inc]  2 [:block [:inc] [:inc]]
                -1 [:dec] -2 [:block [:dec] [:dec]]
                (where val (= (type val) :number)) [:block [:clc] [:adc val]]
                (where sym (= (type sym) :string) (self:local-offset sym)) [:block [:clc] [:adc (self:local-offset sym) :s]]
                _ [:block (self:push) (self:compile-expr val) [:clc] [:adc 0 :s] (self:drop)]))))
        block)
   :- (lambda [self first ...]
        (let [block [:block (self:compile-expr first)]
              nargs (select :# ...)]
          (if (= nargs 0) (table.insert block [:block [:eor 0xffff] [:inc]]) ; negate with two's complement
            (for [i 1 nargs]
              (let [val (select i ...)]
                (table.insert block
                  (match val
                     1 [:dec]  2 [:block [:dec] [:dec]]
                    -1 [:inc] -2 [:block [:inc] [:inc]]
                    (where val (= (type val) :number)) [:block [:sec] [:sbc val]]
                    (where sym (= (type sym) :string) (self:local-offset sym)) [:block [:sec] [:sbc (self:local-offset sym) :s]]
                    _ [:block (self:push) (self:compile-expr val) [:sec] [:sbc 0 :s] (self:drop)])))))
          block))
    := (lambda [self lhs rhs] (boolop self lhs rhs :beq))
    :not= (lambda [self lhs rhs] (boolop self lhs rhs :bne))
    :< (lambda [self lhs rhs] (boolop self lhs rhs :bmi))
    :> (lambda [self lhs rhs] (boolop self rhs lhs :bmi))
    :>= (lambda [self lhs rhs] (boolop self lhs rhs :bpl))
    :<= (lambda [self lhs rhs] (boolop self rhs lhs :bpl))
    :not (lambda [self bool] (self:compile-expr [:if bool 0 0xffff]))
    :if (lambda [self test iftrue ?iffalse]
          (lume.concat
            [:block
              (self:compile-expr test)
              [:cmp 0]
              [:beq :-elseblock-]
              (self:compile-expr iftrue)]
              (when ?iffalse [[:bra :-finished-]])
            [:-elseblock-]
              (when ?iffalse [(self:compile-expr ?iffalse)])
            [:-finished-]))
    :fn (lambda [self name args ...]
          (assert (= (length self.locals) 0))
          (set self.locals (lume.concat args [[:tmp]]))
          (tset self.functions name {:arity (length args) : args :org self.org})
          ; todo: maybe handle mutually recursive functions? (compile-expr only has access to currently-defined functions)
          (self.org:append name (self:compile-expr (lume.concat [:do ...] [[:asm [:rts]]])))
          (assert (= (length self.locals) (+ (length args) 1)))
          (set self.locals []))
    :form (lambda [self name func] (tset self.forms name func))
  })

(fn Ssc.local-offset [self symbol]
  (var offset nil)
  (for [i 1 (length self.locals)]
    (when (= (. self.locals i) symbol)
      (set offset (* 2 (- (length self.locals) i)))))
  offset)

(fn Ssc.compile-expr [self expr]
  (match expr
    (where num (= (type num) :number)) [:lda num]
    (where loc (= (type loc) :string) (self:local-offset loc)) [:lda (self:local-offset loc) :s]
    ; TODO: Global scope
    (where [func & args] (= (?. self.functions func :arity) (length args)))
      (let [pre (icollect [_ arg (ipairs args)] [:block (self:compile-expr arg) (self:push)])
            post (icollect [_ (countiter (length args))] (self:drop))]
        (lume.concat [:block] pre [[:jsr func]] post))
    (where [form & args] (. self.forms form))
      ((. self.forms form) self (table.unpack args))
    nil [:block]
    _ (error (.. "Unrecognized expression " (fv expr)))))

(fn Ssc.compile [self ...]
  (for [i 1 (select :# ...)]
    (self:compile-expr (select i ...)))
  self)

Ssc