; 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.

; * Values default to 16-bit integers, like Forth or BCPL.
; * 32-bit integers are also handled, and the "last value" is stored in a separate "register" in the direct page.
; * 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)

(import-macros {:sss ! : compile} :ssc.macros)
(local Object (require :core.object))
(local lume (require :lib.lume))
(local Ssc (Object:extend))
(local Prg (require :asm.asm))
(local util (require :lib.util))

(set Ssc.LONG_LO :d0x00)
(set Ssc.LONG_HI :d0x02)

(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 [])
  (set self.modules {})
  (if opts.boot (self:compile (table.unpack opts.boot))
    (self:compile (!
      (start-symbol boot)
      (org [(or opts.boot-org 0)])
      (asm
        boot
        (clc) (xce) (rep 0x30) ; disable emulation mode, 16-bit index registers
        (jsr main)
        (sec) (xce) ;re-enter emulation mode
        (rts))))))

(fn Ssc.push [self name expr etype]
  (table.insert self.locals {: name :type etype})
  (match etype
    :word [:block expr [:pha]]
    :long [:block expr [:lda self.LONG_HI] [:pha] [:lda self.LONG_LO] [:pha]]))

(fn Ssc.remove-local [self ?name]
  (let [loc (. self.locals (length self.locals))]
    (when (not= loc.name ?name) (error (.. "Internal stack error: expected " (or ?name "temporary") ", got " (or loc.name "temporary"))))
    (tset self.locals (length self.locals) nil)
    loc))

(fn Ssc.drop [self ?name]
  (match (. (self:remove-local ?name) :type)
    :word [:ply]
    :long [:block [:ply] [:ply]]))

(fn Ssc.pop [self ?name]
  (let [{:type etype} (self:remove-local ?name)]
    (values (match etype
              :word [:pla]
              :long [:block [:pla] [:sta self.LONG_LO] [:pla] [:sta self.LONG_HI]])
            etype)))

(fn Ssc.was-dropped [self localcount]
  (set self.locals (lume.slice self.locals 1 (- (length self.locals) localcount))))

(fn Ssc.defining? [self] (> (length self.locals) 0))

(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)))))))

; Comparison theory:
; word x word constant -> word
; long x long constant -> word
; word x word -> word
; long x long -> word
; Any combination of word and long gets promoted to long (we have to compare both words anyway).

; this should hold for beq and bpl, but not bmi and bne
; optimized cases:
; reg + immediate
; reg + stack? nope, reg + stack isn't really possible when the operator isn't commutative
; stack + reg

(fn boolop [self left right branch]
  (let [compiled-left    (self:expr-word left)
        push-left        (when (not= (type right) :number) (self:push))
        compiled-compare (if (not push-left) [:cmp right]
                           [:block push-left (self:expr-word right) [:cmp 1 :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]))

; operations that work on the accumulator, like adc or sbc
; optimization strategy: keep the current result in the accumulator, work from the stack or immediate values
; 1. take "right" arguments and push them (unless already on stack, immediate, or absolute)
; 2. load left into accumulator
; 3. apply until done
(fn Ssc.accumulation-op [self op first ...]
  (var etype (self:type-expr first))
  (for [i 1 (select :# ...)] (when (= (self:type-expr (select i ...)) :long) (set etype :long)))
  (let [args (icollect [_ val (ipairs [...])] (self:push-addressible val))
        setup (icollect [_ {: setup} (ipairs args)] (when setup (setup)))
        acc (: self (.. :expr- etype) first)
        operations (icollect [i addr (ipairs args)] (op etype addr i))
        cleanup (icollect [_ {: cleanup} (ipairs args)] (when cleanup (cleanup)))]
    (values (lume.concat [:block] setup [acc] operations cleanup) etype)))

(set Ssc.forms
  {:asm (fn [self ...]
          (if (self:defining?) [:block ...]
            (self.org:append ...)))
   :org (lambda [self org] (set self.org (self.prg:org org)))
   :start-symbol (lambda [self symbol] (set self.prg.start-symbol symbol))
   :form (lambda [self name func] (tset self.forms name func))
   :require (lambda [self name]
              (when (= (. self.modules name) nil)
                (let [mod (util.reload name)
                      func (if (= (type mod) :function) mod mod.module)]
                  (tset self.modules name mod)
                  (func self))))
   :do (fn [self ...]
         (var etype-body :void)
         (local c-body (lume.concat [:block] (icollect [i (countiter (select :# ...))]
                                               (let [(expr etype) (self:expr-poly (select i ...))]
                                                 (set etype-body etype)
                                                 expr))))
         (values c-body etype-body))
   :let (fn [self bindings ...]
          (let [compiled-bindings    (icollect [_ symbol expr (pairoff bindings)] (self:push symbol (self:expr-poly expr)))
                (compiled-body etype) (self:expr-poly [:do ...])
                compiled-cleanup     (icollect [i-half (countiter (/ (length bindings) 2))]
                                       (self:drop (. bindings (- (length bindings) (* i-half 2) -1))))]
            (values (lume.concat [:block] compiled-bindings [compiled-body] compiled-cleanup) etype)))
   :fn (lambda [self name args ...]
         (assert (not (self:defining?)) "Can't nest function definitions")
         (local arglocals (icollect [_ arg (ipairs args)] (match arg
                                                           [:long aname] {:name aname :type :long}
                                                           aname         {:name aname :type :word})))
         (set self.locals (lume.concat arglocals [{:type :word :comment :returnaddr}]))
         ; todo: maybe handle mutually recursive functions? (compile-expr only has access to currently-defined functions)
         (local (c-function etype) (self:expr-poly [:do ...]))
         (self.org:append name c-function [:rts])
         (tset self.functions name {:arity (length args) :args arglocals :org self.org :type etype})
         (assert (= (length self.locals) (+ (length args) 1))
                 (.. "Left locals on stack?? Expected " (tostring (+ (length args) 1)) " got " (tostring (length self.locals))))
         (set self.locals []))
   :if (lambda [self test iftrue ?iffalse]
         (let [c-test (self:expr-word test)
               (c-true truetype) (self:expr-poly iftrue)
               (c-false falsetype) (when ?iffalse (self:expr-poly ?iffalse))
               etype (if (and falsetype (not= truetype falsetype)) :void truetype)]
           (values (lume.concat [:block
                                   c-test
                                   [:cmp 0] [:beq (if ?iffalse :-elseblock- :-finished-)]
                                   c-true]
                                (when ?iffalse [
                                    [:bra :-finished-]
                                  :-elseblock- c-false])
                                [:-finished-]) etype)))

   :+ (lambda [self first ...]
        (self:accumulation-op
          (fn [etype {: lo : hi : const}]
            (match etype
              :word (if (= const 1)  [:inc] (= const 2)  [:block [:inc] [:inc]]
                        (= const -1) [:dec] (= const -2) [:block [:dec] [:dec]]
                      [:block [:clc] (lo :adc)])
              :long [:block [:clc] (lo :lda) [:adc self.LONG_LO] [:sta self.LONG_LO]
                                   (if hi (hi :lda) [:lda 0]) [:adc self.LONG_HI] [:sta self.LONG_HI]]))
          first ...))
   :- (lambda [self first ...]
        (self:accumulation-op
          (fn [etype {: lo : hi : const}]
            (match etype
              :word (if (= const 1)  [:dec] (= const 2)  [:block [:dec] [:dec]]
                        (= const -1) [:inc] (= const -2) [:block [:inc] [:inc]]
                      [:block [:sec] (lo :sbc)])
              :long [:block [:sec] (lo :lda) [:sbc self.LONG_LO] [:sta self.LONG_LO]
                                   (if hi (hi :lda) [:lda 0]) [:sbc self.LONG_HI] [:sta self.LONG_HI]]))))
    := (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:expr-poly [:if bool 0 0xffff]))
    :or (lambda [self lhs rhs] (self:expr-poly [:if lhs 0xffff [:if rhs 0xffff 0]]))
    :and (lambda [self lhs rhs] (self:expr-poly [:if lhs [:if rhs 0xffff 0] 0]))
    :& (lambda [self label] [:lda #($1:lookup-addr label)])
  })

(fn Ssc.local-offset [self name-or-index]
  (var offset nil)
  (for [i 1 (length self.locals)]
    (when (or (= i name-or-index) (= (. self.locals i :name) name-or-index))
      (set offset (+ 1 (* 2 (- (length self.locals) i))))))
  offset)

(fn Ssc.local-type [self name-or-index]
  (var etype nil)
  (for [i 1 (length self.locals)]
    (when (or (= i name-or-index) (= (. self.locals i :name) name-or-index))
      (set etype (. self.locals i :type))))
  etype)

(fn Ssc.set-long [self loexpr hiexpr]
  [:block loexpr [:sta self.LONG_LO] hiexpr [:sta self.LONG_HI]])
(fn Ssc.set-long-local [self loc]
  (self:set-long [:lda (self:local-offset loc) :s] [:lda (+ (self:local-offset loc) 2) :s]))

(fn Ssc.type-expr [self expr] (let [(_ etype) (self:expr-poly expr)] etype))

(fn Ssc.addressible-const [self const]
  {:lo #[$1 (bit.band const 0xffff)] :hi (let [hi (bit.rshift (bit.band 0xffff0000) 16)] (if (or (= hi 0) (= hi 0xffff)) nil #[$1 hi])) : const})
(fn Ssc.addressible-loc [self loc]
  {:lo #[$1 (self:local-offset loc) :s] :hi (when (= (self:local-type loc) :long) #[$1 (+ (self:local-offset loc) 2)])})

(fn Ssc.addressible [self expr]
  (match (type expr)
    :number                                  (self:addressible-const expr)
    (where :string (self:local-offset expr)) (self:addressible-loc expr)))

(fn Ssc.push-addressible [self expr]
  (or (self:addressible expr)
      (let [c (self:push nil (self:expr-poly expr))
            iloc (length (self.locals))]
        (lume.merge (self:addressible-loc iloc) {:setup #c :cleanup #(self:drop)}))))

(fn Ssc.expr-poly [self expr]
  (match expr
    (where lit (?. (self:addressible lit) :hi)) (let [{: lo : hi} (self:addressible lit)]
      (values [:block (lo :lda) [:sta self.LONG_LO] (hi :lda) [:sta self.LONG_HI]] :long))
    (where lit (?. (self:addressible lit) :lo)) (let [{: lo} (self:addressible lit)] (values (lo :lda) :word))
    ; TODO: Global scope
    (where [func & args] (= (?. self.functions func :arity) (length args)))
      (let [f (. self.functions func)
            pre (icollect [iarg arg (ipairs args)]
                  (let [atype (. f.args iarg :type)
                        c-arg (: self (.. :expr- atype) arg)
                        c-push (self:push nil c-arg atype)]
                    c-push))
            post (icollect [_ (countiter (length args))] (self:drop))]
        (print (fv pre) (fv post) (fv args))
        (values (lume.concat [:block] pre [[:jsr func]] post) f.type))
    (where [form & args] (. self.forms form))
      (let [f (. self.forms form)
            (cexpr etype) (f self (table.unpack args))]
        (values cexpr (or etype :word)))
    nil (values [:block] :void)
    _ (error (.. "Unrecognized expression " (fv expr)))))

(fn Ssc.expr-word [self expr]
  (let [(c etype) (self:expr-poly expr)]
    (when (not= etype :word) (error "Unexpected long or void in " (fv expr) " - please wrap in explicit truncation form"))
    c))

(fn Ssc.expr-long [self expr]
  (let [(c etype) (self:expr-poly expr)]
    (match etype
      :long c
      :word [:block c [:sta self.LONG_LO] [:lda 0] [:sta self.LONG_HI]]
      _ (error "Unexpected type " etype " in " (fv expr) " - wanted long or word"))))

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

(fn Ssc.assemble [self]
  (self.prg:assemble)
  self.prg)

Ssc