;; Forth runtime — state, stacks, dictionary, output buffer.
;; Data stack: mutable SX list, TOS = first.
;; Return stack: separate mutable list.
;; Dictionary: SX dict {lowercased-name -> word-record}.
;; Word record: {"kind" "body" "immediate?"}; kind is "primitive" or "colon-def".
;; Output buffer: mutable string appended to by `.`, `EMIT`, `CR`, etc.
;; Compile-mode flag: "compiling" on the state.

(define
  forth-make-state
  (fn
    ()
    (let
      ((s (dict)))
      (dict-set! s "dstack" (list))
      (dict-set! s "rstack" (list))
      (dict-set! s "dict" (dict))
      (dict-set! s "output" "")
      (dict-set! s "compiling" false)
      (dict-set! s "current-def" nil)
      (dict-set! s "base" 10)
      (dict-set! s "vars" (dict))
      s)))

(define
  forth-error
  (fn (state msg) (dict-set! state "error" msg) (raise msg)))

(define
  forth-push
  (fn (state v) (dict-set! state "dstack" (cons v (get state "dstack")))))

(define
  forth-pop
  (fn
    (state)
    (let
      ((st (get state "dstack")))
      (if
        (= (len st) 0)
        (forth-error state "stack underflow")
        (let ((top (first st))) (dict-set! state "dstack" (rest st)) top)))))

(define
  forth-peek
  (fn
    (state)
    (let
      ((st (get state "dstack")))
      (if (= (len st) 0) (forth-error state "stack underflow") (first st)))))

(define forth-depth (fn (state) (len (get state "dstack"))))

(define
  forth-rpush
  (fn (state v) (dict-set! state "rstack" (cons v (get state "rstack")))))

(define
  forth-rpop
  (fn
    (state)
    (let
      ((st (get state "rstack")))
      (if
        (= (len st) 0)
        (forth-error state "return stack underflow")
        (let ((top (first st))) (dict-set! state "rstack" (rest st)) top)))))

(define
  forth-rpeek
  (fn
    (state)
    (let
      ((st (get state "rstack")))
      (if
        (= (len st) 0)
        (forth-error state "return stack underflow")
        (first st)))))

(define
  forth-emit-str
  (fn (state s) (dict-set! state "output" (str (get state "output") s))))

(define
  forth-make-word
  (fn
    (kind body immediate?)
    (let
      ((w (dict)))
      (dict-set! w "kind" kind)
      (dict-set! w "body" body)
      (dict-set! w "immediate?" immediate?)
      w)))

(define
  forth-def-prim!
  (fn
    (state name body)
    (dict-set!
      (get state "dict")
      (downcase name)
      (forth-make-word "primitive" body false))))

(define
  forth-def-prim-imm!
  (fn
    (state name body)
    (dict-set!
      (get state "dict")
      (downcase name)
      (forth-make-word "primitive" body true))))

(define
  forth-lookup
  (fn (state name) (get (get state "dict") (downcase name))))

(define
  forth-binop
  (fn
    (op)
    (fn
      (state)
      (let
        ((b (forth-pop state)) (a (forth-pop state)))
        (forth-push state (op a b))))))

(define
  forth-unop
  (fn
    (op)
    (fn (state) (let ((a (forth-pop state))) (forth-push state (op a))))))

(define
  forth-cmp
  (fn
    (op)
    (fn
      (state)
      (let
        ((b (forth-pop state)) (a (forth-pop state)))
        (forth-push state (if (op a b) -1 0))))))

(define
  forth-cmp0
  (fn
    (op)
    (fn
      (state)
      (let ((a (forth-pop state))) (forth-push state (if (op a) -1 0))))))

(define
  forth-trunc
  (fn (x) (if (< x 0) (- 0 (floor (- 0 x))) (floor x))))

(define
  forth-div
  (fn
    (a b)
    (if (= b 0) (raise "division by zero") (forth-trunc (/ a b)))))

(define
  forth-mod
  (fn
    (a b)
    (if (= b 0) (raise "division by zero") (- a (* b (forth-div a b))))))

(define forth-bits-width 32)

(define
  forth-to-unsigned
  (fn (n w) (let ((m (pow 2 w))) (mod (+ (mod n m) m) m))))

(define
  forth-from-unsigned
  (fn
    (n w)
    (let ((half (pow 2 (- w 1)))) (if (>= n half) (- n (pow 2 w)) n))))

(define
  forth-bitwise-step
  (fn
    (op ua ub out place i w)
    (if
      (>= i w)
      out
      (let
        ((da (mod ua 2)) (db (mod ub 2)))
        (forth-bitwise-step
          op
          (floor (/ ua 2))
          (floor (/ ub 2))
          (+ out (* place (op da db)))
          (* place 2)
          (+ i 1)
          w)))))

(define
  forth-bitwise-uu
  (fn
    (op)
    (fn
      (a b)
      (let
        ((ua (forth-to-unsigned a forth-bits-width))
          (ub (forth-to-unsigned b forth-bits-width)))
        (forth-from-unsigned
          (forth-bitwise-step op ua ub 0 1 0 forth-bits-width)
          forth-bits-width)))))

(define
  forth-bit-and
  (forth-bitwise-uu (fn (x y) (if (and (= x 1) (= y 1)) 1 0))))

(define
  forth-bit-or
  (forth-bitwise-uu (fn (x y) (if (or (= x 1) (= y 1)) 1 0))))

(define forth-bit-xor (forth-bitwise-uu (fn (x y) (if (= x y) 0 1))))

(define forth-bit-invert (fn (a) (- 0 (+ a 1))))

(define
  forth-install-primitives!
  (fn
    (state)
    (forth-def-prim! state "DUP" (fn (s) (forth-push s (forth-peek s))))
    (forth-def-prim! state "DROP" (fn (s) (forth-pop s)))
    (forth-def-prim!
      state
      "SWAP"
      (fn
        (s)
        (let
          ((b (forth-pop s)) (a (forth-pop s)))
          (forth-push s b)
          (forth-push s a))))
    (forth-def-prim!
      state
      "OVER"
      (fn
        (s)
        (let
          ((b (forth-pop s)) (a (forth-pop s)))
          (forth-push s a)
          (forth-push s b)
          (forth-push s a))))
    (forth-def-prim!
      state
      "ROT"
      (fn
        (s)
        (let
          ((c (forth-pop s)) (b (forth-pop s)) (a (forth-pop s)))
          (forth-push s b)
          (forth-push s c)
          (forth-push s a))))
    (forth-def-prim!
      state
      "-ROT"
      (fn
        (s)
        (let
          ((c (forth-pop s)) (b (forth-pop s)) (a (forth-pop s)))
          (forth-push s c)
          (forth-push s a)
          (forth-push s b))))
    (forth-def-prim!
      state
      "NIP"
      (fn (s) (let ((b (forth-pop s))) (forth-pop s) (forth-push s b))))
    (forth-def-prim!
      state
      "TUCK"
      (fn
        (s)
        (let
          ((b (forth-pop s)) (a (forth-pop s)))
          (forth-push s b)
          (forth-push s a)
          (forth-push s b))))
    (forth-def-prim!
      state
      "?DUP"
      (fn
        (s)
        (let ((a (forth-peek s))) (when (not (= a 0)) (forth-push s a)))))
    (forth-def-prim! state "DEPTH" (fn (s) (forth-push s (forth-depth s))))
    (forth-def-prim!
      state
      "PICK"
      (fn
        (s)
        (let
          ((n (forth-pop s)) (st (get s "dstack")))
          (if
            (or (< n 0) (>= n (len st)))
            (forth-error s "PICK out of range")
            (forth-push s (nth st n))))))
    (forth-def-prim!
      state
      "ROLL"
      (fn
        (s)
        (let
          ((n (forth-pop s)) (st (get s "dstack")))
          (if
            (or (< n 0) (>= n (len st)))
            (forth-error s "ROLL out of range")
            (let
              ((taken (nth st n))
                (before (take st n))
                (after (drop st (+ n 1))))
              (dict-set! s "dstack" (concat before after))
              (forth-push s taken))))))
    (forth-def-prim!
      state
      "2DUP"
      (fn
        (s)
        (let
          ((b (forth-pop s)) (a (forth-pop s)))
          (forth-push s a)
          (forth-push s b)
          (forth-push s a)
          (forth-push s b))))
    (forth-def-prim! state "2DROP" (fn (s) (forth-pop s) (forth-pop s)))
    (forth-def-prim!
      state
      "2SWAP"
      (fn
        (s)
        (let
          ((d (forth-pop s))
            (c (forth-pop s))
            (b (forth-pop s))
            (a (forth-pop s)))
          (forth-push s c)
          (forth-push s d)
          (forth-push s a)
          (forth-push s b))))
    (forth-def-prim!
      state
      "2OVER"
      (fn
        (s)
        (let
          ((d (forth-pop s))
            (c (forth-pop s))
            (b (forth-pop s))
            (a (forth-pop s)))
          (forth-push s a)
          (forth-push s b)
          (forth-push s c)
          (forth-push s d)
          (forth-push s a)
          (forth-push s b))))
    (forth-def-prim! state "+" (forth-binop (fn (a b) (+ a b))))
    (forth-def-prim! state "-" (forth-binop (fn (a b) (- a b))))
    (forth-def-prim! state "*" (forth-binop (fn (a b) (* a b))))
    (forth-def-prim! state "/" (forth-binop forth-div))
    (forth-def-prim! state "MOD" (forth-binop forth-mod))
    (forth-def-prim!
      state
      "/MOD"
      (fn
        (s)
        (let
          ((b (forth-pop s)) (a (forth-pop s)))
          (forth-push s (forth-mod a b))
          (forth-push s (forth-div a b)))))
    (forth-def-prim! state "NEGATE" (forth-unop (fn (a) (- 0 a))))
    (forth-def-prim! state "ABS" (forth-unop abs))
    (forth-def-prim!
      state
      "MIN"
      (forth-binop (fn (a b) (if (< a b) a b))))
    (forth-def-prim!
      state
      "MAX"
      (forth-binop (fn (a b) (if (> a b) a b))))
    (forth-def-prim! state "1+" (forth-unop (fn (a) (+ a 1))))
    (forth-def-prim! state "1-" (forth-unop (fn (a) (- a 1))))
    (forth-def-prim! state "2+" (forth-unop (fn (a) (+ a 2))))
    (forth-def-prim! state "2-" (forth-unop (fn (a) (- a 2))))
    (forth-def-prim! state "2*" (forth-unop (fn (a) (* a 2))))
    (forth-def-prim! state "2/" (forth-unop (fn (a) (floor (/ a 2)))))
    (forth-def-prim! state "=" (forth-cmp (fn (a b) (= a b))))
    (forth-def-prim! state "<>" (forth-cmp (fn (a b) (not (= a b)))))
    (forth-def-prim! state "<" (forth-cmp (fn (a b) (< a b))))
    (forth-def-prim! state ">" (forth-cmp (fn (a b) (> a b))))
    (forth-def-prim! state "<=" (forth-cmp (fn (a b) (<= a b))))
    (forth-def-prim! state ">=" (forth-cmp (fn (a b) (>= a b))))
    (forth-def-prim! state "0=" (forth-cmp0 (fn (a) (= a 0))))
    (forth-def-prim! state "0<>" (forth-cmp0 (fn (a) (not (= a 0)))))
    (forth-def-prim! state "0<" (forth-cmp0 (fn (a) (< a 0))))
    (forth-def-prim! state "0>" (forth-cmp0 (fn (a) (> a 0))))
    (forth-def-prim! state "AND" (forth-binop forth-bit-and))
    (forth-def-prim! state "OR" (forth-binop forth-bit-or))
    (forth-def-prim! state "XOR" (forth-binop forth-bit-xor))
    (forth-def-prim! state "INVERT" (forth-unop forth-bit-invert))
    (forth-def-prim!
      state
      "."
      (fn (s) (forth-emit-str s (str (forth-pop s) " "))))
    (forth-def-prim!
      state
      ".S"
      (fn
        (s)
        (let
          ((st (reverse (get s "dstack"))))
          (forth-emit-str s "<")
          (forth-emit-str s (str (len st)))
          (forth-emit-str s "> ")
          (for-each (fn (v) (forth-emit-str s (str v " "))) st))))
    (forth-def-prim!
      state
      "EMIT"
      (fn (s) (forth-emit-str s (code-char (forth-pop s)))))
    (forth-def-prim! state "CR" (fn (s) (forth-emit-str s "\n")))
    (forth-def-prim! state "SPACE" (fn (s) (forth-emit-str s " ")))
    (forth-def-prim!
      state
      "SPACES"
      (fn
        (s)
        (let
          ((n (forth-pop s)))
          (when
            (> n 0)
            (for-each (fn (_) (forth-emit-str s " ")) (range 0 n))))))
    (forth-def-prim! state "BL" (fn (s) (forth-push s 32)))
    state))