; hotswap support
(local {: addr-parser} (require :asm.65816))
(local util (require :lib.util))
(local {: bytes-to-uint16 : bytes-to-uint24 : int16-to-bytes : int24-to-bytes} util)

; Hotswap theory:
; The common case is code moving around in memory; even if you are not changing the content of code that
; is currently executing, any changes to anything mean that pointers to that code will need to change.
; If you try not to store pointers in globals (when everything is statically allocated, you should explicitly
; name things, like in Forth) then the main place that persistent pointers to code exist is in call stacks -
; specifically, the values on the stack consumed by "rts" or "rtl". When hotswapping, we need to walk the
; callstack and A. find these values, and B. patch them to their new values. For this, we need to be able to
; map any address that might be returned to from the old program to the new program - we can easily do this
; by generating symbols for each callsite and setting up a reverse lookup table. We _also_ need to know the
; stack layout at the time of each call so that we can find the next link in the chain, and to verify that
; this hasn't changed (if the function has changed enough, then we can't modify it mid-execution).

(fn assert-local-matches [funcname loc-old loc-new]
  (each [_ key (ipairs [:type :name :returnaddr])]
    (assert (= (. loc-old key) (. loc-new key)) (.. "Stack mismatch when patching " funcname))))

(fn assert-locals-match [funcname locals-old locals-new]
  (assert (= (length locals-old) (length locals-new)) (.. "Stack size mismatch when patching " funcname))
  (each [iloc loc-old (ipairs locals-old)]
    (assert-local-matches funcname loc-old (. locals-new iloc))))

(fn next-callsite-offset [locals]
  (var offset 0)
  (var start-counting false)
  (each [_ loc (ipairs locals)]
    (if start-counting (let [size (match loc.type :placeholder 0 :word 2 :long 4)] (set offset (+ offset size)))
        loc.returnaddr (set start-counting true)))
  offset)

(fn next-callsite-far? [locals]
  (var far? false)
  (each [_ loc (ipairs locals)]
    (when (and loc.returnaddr (= loc.type :long))
      (set far? true)))
  far?)

(fn read-callsite-addr [stack bank far]
  (if far (bytes-to-uint24 stack 1) ; lowest byte (top of stack) is preserved B register
          (bit.bor (bytes-to-uint16 stack) (bit.lshift bank 16))))

(fn lookup-callsite [ssc addr]
  (if (= addr (- (ssc.prg:lookup-addr :yield-forever) 1))
      {:callsite-sym :yield-forever :locals [] :calling :yield-forever :funcname "<base task>"}
    (. ssc.addr-to-callsite addr)))

(fn patch-stack [ssc-old ssc-new stack bank far]
  (print (length stack) bank far)
  (if (= (length stack) 0) stack
      ; top-of-stack should be a callsite; look it up
    (let [callsite-addr (read-callsite-addr stack bank far)
          {: callsite-sym : locals : funcname} (assert (lookup-callsite ssc-old callsite-addr)
                                                 (.. "Top of stack value " callsite-addr " is not a recognized callsite"))
          new-addr (- (ssc-new.prg:lookup-addr callsite-sym) 1)
          {:locals new-locals} (lookup-callsite ssc-new  new-addr)
          new-bank (bit.rshift new-addr 16)
          _ (when (not far) (assert (= (bit.band callsite-addr 0xff0000) (bit.band new-addr 0xff0000))
                              (.. funcname " moved banks from " bank " to " new-bank)))
          _ (print (.. "patching " callsite-sym " from " callsite-addr " to " new-addr))
          _ (assert-locals-match funcname locals new-locals)
          new-top (if far (.. (stack:sub 1 1) (int24-to-bytes new-addr))
                          (int16-to-bytes new-addr))
          iaftertop (if far 5 3)
          inextstack (+ iaftertop (next-callsite-offset locals))
          _ (print (.. "locals " (fv locals) " offset " (next-callsite-offset locals)) inextstack)]
      (.. new-top (stack:sub iaftertop (- inextstack 1))
          (if (= funcname ssc-old.prg.start-symbol) (stack:sub inextstack) ; stop when we hit the boot-up function
            (patch-stack ssc-old ssc-new (stack:sub inextstack) new-bank (next-callsite-far? locals)))))))

(fn split-equally [s size]
  (values (fn [s iprev] (let [i (+ iprev 1)
                              istart (+ (* (- i 1) size) 1)
                              iend (* i size)]
                          (when (>= (length s) iend) (values i (s:sub istart iend)))))
          s 0))

(fn read-stacks [link ssc]
  (let [stack-bounds-addr (ssc.prg:lookup-addr :first-task)
        stack-bounds-bytes (link:read stack-bounds-addr 4)
        first-task (bytes-to-uint16 stack-bounds-bytes)
        last-task (bytes-to-uint16 stack-bounds-bytes 2)
        task-size 0x100
        read-size (+ (- last-task first-task) task-size)
        task-count (/ read-size task-size)
        task-bytes (link:read first-task read-size)
        sp-offset (addr-parser ssc.TASK-STACK)]
    (icollect [_ task (split-equally task-bytes task-size)]
      (let [sp-addr (bytes-to-uint16 task sp-offset)
            istackstart (bit.band sp-addr 0xff)
            stack (task:sub (+ istackstart 2))]
        {:sp-addr (+ sp-addr 1) : stack}))))

(fn lookup-yield-bank [ssc] (bit.rshift (ssc.prg:lookup-addr :yield) 16))
(fn filter-nonyielding-stacks [ssc stacks]
  (let [yield-bank (lookup-yield-bank ssc)]
    (print "yield-bank" yield-bank)
    (icollect [_ stack (ipairs stacks)]
      (let [callsite-addr (read-callsite-addr stack.stack yield-bank false)
            callsite (lookup-callsite ssc callsite-addr)
            funcname (?. callsite :calling)]
        (print "top callsite:" (string.format "%X" callsite-addr) (fv callsite))
        (when (= funcname :yield) stack)))))

(fn hotswap-stacks [link ssc-old ssc-new]
  (print "hotswappin time" link ssc-old ssc-new)
  (let [stacks (read-stacks link ssc-old)
        _ (print "read stacks" (fv stacks))
        yielding-stacks (filter-nonyielding-stacks ssc-old stacks)]
    (print "filtered stacks" (length stacks) (length yielding-stacks))
    (collect [_ {: sp-addr : stack} (ipairs yielding-stacks)]
      (let [(success new-stack) (pcall #(patch-stack ssc-old ssc-new stack (lookup-yield-bank ssc-old) false))]
        (if success (do (print "patched" (stack:tohex) "to" (new-stack:tohex)) (values (tostring sp-addr) new-stack))
          (error (.. new-stack ": stack at " (string.format "%X" sp-addr))))))))

{: hotswap-stacks}