; Lisp compiler to ARM Thumb 2 Assembler - Version 3 - 8th September 2025
;

#|
Language definition:

Defining variables and functions: defun, setq
Symbols: nil, t
List functions: car, first, cdr, rest, nth
Unary arithmetic functions: -, 1+, 1-
Binary arithmetic functions: +, -, *, /, mod, logand, logior, logxor, ash
Unary predicates: zerop, plusp, minusp, oddp, evenp
Binary predicates: =, <, <=, >, >=, /=
Unary logical function: not
Binary logical functions: and, or
Conditionals: if, when, unless
Evaluation: progn, let, let*
Iteration: dotimes, loop, return
|#

; Compile a lisp function

(defun compile (name)
  (if (eq (car (eval name)) 'lambda)
      (eval (comp (cons 'defun (cons name (cdr (eval name))))))
    (error "Not a Lisp function")))

; The main compile routine - returns compiled code for x, prefixed by type :integer or :boolean
; Leaves result in r0

(defun comp (x &optional env)
  (cond
   ((null x) (type-code :boolean '(($mov 'r0 0))))
   ((eq x t) (type-code :boolean '(($mov 'r0 1))))
   ((symbolp x) (comp-symbol x env))
   ((integerp x) (comp-integer x))
   (t (let ((fn (first x)) (args (rest x)))
        (case fn
          (defun   (setq *label-num* 0)
                   (setq *used* (subseq *locals* 0 (length (second args))))
                   (setq env (mapcar #'(lambda (x y) (cons x y)) (second args) *used*))
                   (comp-defun (first args) (second args) (cddr args) env))
          (setq    (comp-setq (first args) (second args) env))
          (let     (comp-let (first args) (rest args) nil env))
          (let*    (comp-let (first args) (rest args) t env))
          (progn   (comp-progn args env))
          (if      (comp-if (first args) (list (second args)) (list (third args)) env))
          (when    (comp-if (first args) (rest args) nil env))
          (unless  (comp-if (first args) nil (rest args) env))
          (dotimes (comp-dotimes (first args) (rest args) env))
          (loop    (comp-loop args env))
          (return  (comp-return (first args) env))
          (t       (comp-funcall fn args env)))))))

; Utilities

; Thumb-2 instructions

(defun cbz-cbnz (op argn label)
  (let* ((rn (regno argn))
         (soff6 (logand (offset label) #x3f))
         (i (ash soff6 -5))
         (immed5 (logand soff6 #x1f)))
    (emit #x61153000 op i 1 immed5 rn)))

(defun $cbz (argn label)
  (cbz-cbnz 44 argn label))

(defun $cbnz (argn label)
  (cbz-cbnz 46 argn label))

(defun condcode (cc)
  (case cc (eq 0) (ne 1) ((cs hs) 2) ((cc lo) 3) (mi 4) (pl 5) (vs 6) (vc 7)
    (hi 8) (ls 9) (ge 10) (lt 11) (gt 12) (le 13) (al 14)))

(defun makemask (xyz fc0)
  (cond
   ((zerop (length xyz)) #b1000)
   (t (logior
       (ash (if (eq (char xyz 0) #\t) fc0 (- 1 fc0)) 3)
       (ash (makemask (subseq xyz 1) fc0) -1)))))

(defun $it (xyz firstcond)
  (let* ((cc (condcode firstcond))
        (fc0 (logand cc 1)))
    (emit #x84400000 #xbf cc (makemask (string xyz) fc0))))

(defun $sdiv (argd argn &optional argm)
  (unless argm (setq argm argn argn argd))
  (let ((rd (regno argd)) (rn (regno argn)) (rm (regno argm)))
    (list 
     (emit #xc4000000 #xfb9 rn)
     (emit #x44440000 #xf rd #xf rm))))

(defun $mls (argd argn argm arga)
  (let ((rd (regno argd)) (rn (regno argn)) (rm (regno argm)) (ra (regno arga)))
    (list 
     (emit #xc4000000 #xfb0 rn)
     (emit #x44440000 ra rd 1 rm))))

(defun movw-movt (op argd imm16)
  (let ((rd (regno argd)) (i (logand (ash imm16 -11) 1)) (imm4 (logand (ash imm16 -12) #xf))
        (imm3 (logand (ash imm16 -8) #x7)) (imm8 (logand imm16 #xff)))
    (list 
     (emit #x51640000 #x1e i op imm4)
     (emit #x13480000 0 imm3 rd imm8))))

(defun $movw (argd imm16)
  (movw-movt #x24 argd imm16))

(defun $movt (argd imm16)
  (movw-movt #x2C argd imm16))

; Like mapcon but not destructive

(defun mappend (fn lst)
  (apply #'append (mapcar fn lst)))

; The type is prefixed onto the list of assembler code instructions

(defun type-code (type code) (cons type code))

(defun code-type (type-code) (car type-code))

(defun code (type-code) (cdr type-code))

(defun checktype (fn type check)
  (unless (or (null type) (null check) (eq type check))
    (format t "Warning: Argument to '~a' should be ~a not ~a~%" fn type check)))

(defun reg-op (op arg1 arg2)
  (list op (if (symbolp arg1) (list 'quote arg1) arg1)
       (if (symbolp arg2) (list 'quote arg2) arg2)))

; Constants

(defvar *popr1* '($pop '(r1)))

; Allocate registers

(defvar *params* '(r0 r1 r2 r3))

(defvar *locals* '(r4 r5 r6 r7))

(defvar *available* '(r2 r3 r4 r5 r6 r7)) ; for local variables

(defvar *used* nil)

(defun unused-reg (env) 
  (let ((available *available*))
    (dolist (x available (error "No free registers available ~a" env))
      (unless (member x env :test #'(lambda (a b) (eq a (cdr b))))
        (unless (member x *used*) (push x *used*))
        (return x)))))

; Generate a label

(defvar *label-num* 0)

(defun gen-label ()
  (read-from-string (format nil "label~d" (incf *label-num*))))

; Subfunctions

(defun comp-symbol (x env)
  (let ((reg (cdr (assoc x env))))
    (type-code nil (list (reg-op '$mov 'r0 reg)))))

(defun comp-integer (x)
  (type-code
   :integer
   (cond
    ((<= 0 x 255) (list (list '$mov ''r0 x)))
    ((<= 0 x #xffff) (list (list '$movw ''r0 x)))
    ((zerop (logand x #xffff)) (list (list '$movt ''r0 (logand (ash x -16) #xffff))))
    ((<= -255 x -1) (list (list '$mov ''r0 (- x)) '($neg 'r0 'r0)))
    (t (list (list '$movw ''r0 (logand x #xffff)) (list '$movt ''r0 (logand (ash x -16) #xffff)))))))

(defun comp-defun (name args body env)
  (unless (symbolp name) (error "Invalid function name ~a" name))
  (let ((moves (mapcar #'(lambda (x y) (list (reg-op '$mov x y))) *used* *params*))
        (main (comp-progn body env)))
    (append 
     (list 'defcode name args)
     (list name (list '$push (list 'quote (cons 'lr *used*))))
     (apply #'append moves)
     (code main)
     (list (list '$pop (list 'quote (cons 'pc *used*)))))))

(defun comp-setq (var expr env)
  (let ((value (comp expr env))
        (reg (cdr (assoc var env))))
    (type-code 
     (code-type value) 
     (append (code value) (list (reg-op '$mov reg 'r0))))))

(defun comp-let (args body star env)
  (let* ((env2 env)
         (assigns
          (mappend
           #'(lambda (x) 
               (let ((reg (unused-reg env2)))
                 (push (cons (first x) reg) env2)
                 (when star (setq env env2))
                 (append
                  (code (comp (second x) env))
                  (list (reg-op '$mov reg 'r0)))))
           args))
         (body2 (comp-progn body env2)))
    (type-code (code-type body2) (append assigns (code body2)))))

(defun comp-progn (exps env)
  (let* ((len (1- (length exps)))
         (nlast (subseq exps 0 len))
         (last1 (nth len exps))
         (start (mappend #'(lambda (x) (append (code (comp x env)))) nlast))
         (end (comp last1 env)))
    (type-code (code-type end) (append start (code end)))))

(defun comp-if (pred then else env)
  (let ((lab1 (gen-label))
        (lab2 (gen-label))
        (test (comp pred env))
        (op (if then '$cbz '$cbnz)))
    (checktype 'if (car test) :boolean)
    (type-code :integer
               (append
                (code test)
                (list (list op ''r0 lab1))
                (when then (code (comp-progn then env)))
                (when (and then else) (list (list '$b lab2) lab1))
                (when else (code (comp-progn else env)))
                (if (and then else) (list lab2) (list lab1))))))

(defun comp-dotimes (args body env)
  (let* ((var (first args))
         (val (second args))
         (reg (unused-reg env))
         (limit (comp val env))
         (lab1 (gen-label))
         (lab2 (gen-label)))
    (checktype 'dotimes (code-type limit) :integer)
    (push (cons var reg) env)
    (push (cons 'return lab2) env)
    (let* ((body (comp-progn body env))
           (result (third args))
           (main
            (append (list (reg-op '$mov reg 0)) (code limit) 
                    (list lab1 '($push '(r0))) (code body) 
                    (list (reg-op '$add reg 1) '($pop '(r0)) (reg-op '$cmp reg ''r0) (list '$bne lab1)))))
          (let ((return (comp result env)))
            (type-code (code-type return) (append main (code return) (list lab2)))))))

(defun comp-loop (args env)
  (let ((lab1 (gen-label))
        (lab2 (gen-label)))
    (push (cons 'return lab2) env)
    (let ((body (comp-progn args env)))
      (type-code (code-type body)
                 (append
                  (list '($push '(r0)) lab1) ; Allow for pop in comp-return
                  (code body)
                  (list (list '$b lab1) lab2))))))

(defun comp-return (arg env)
  (let ((lab (cdr (assoc 'return env))))
    (unless lab (error "No block to return from"))
    (type-code nil (append (code (comp arg env)) (list *popr1* (list '$b lab))))))

(defun do-predicate (condcode)
  (list (list '$it ''e (list 'quote condcode)) '($mov 'r0 1) '($mov 'r0 0)))

(defun comp-funcall (f args env)
  (let* ((match (assoc f '((zerop test1 1 eq) (plusp test1 1 gt) (minusp test1 1 lt) (oddp test1a 1 cs) (evenp test1a 1 cc)
                           (> test2 2 gt) (>= test2 2 ge) (= test2 2 eq) (<= test2 2 le) (< test2 2 lt) (/= test2 2 ne)
                           (not logical1 1) (and logical2 2 $and) (or logical2 2 $orr) (1+ arith1 1 $add) (1- arith1 1 $sub) (+ arith+- 2 $add) 
                           (- arith+- nil $sub) (* arith2 2 $mul) (logand arith2 2 $and) (logior arith2 2 $orr) (logxor arith2 2 $eor)
                           (ash ash 2) (/ div 2) (truncate div 2) (mod mod 2) (car car 1) (first car 1) (cdr cdr 1) (rest cdr 1)
                           (nth nth 2) (register register nil))))
         (category (second match))
         (nargs (third match))
         (opcode (second (cddr match)))
         (types (case category 
                  ((test1 test1a test2) '(:integer :boolean))
                  ((logical1 logical2) '(:boolean :boolean))
                  ((car first) '(:list :integer))
                  ((cdr rest) '(:list :list))
                  (t '(:integer :integer))))
         (argtype (first types))
         (resulttype (second types)))                       
    (type-code
     resulttype
     (append
      (comp-args f args nargs argtype env)
      (case category
        (test1
         (append (list '($cmp 'r0 0)) (do-predicate opcode)))
        (test1a
         (append (list '($asr 'r0 'r0 1)) (do-predicate opcode)))
        (test2
         (append (list *popr1* '($cmp 'r1 'r0)) (do-predicate opcode)))
        (logical1
         (list '($mov 'r1 1) '($eor 'r0 'r1))) ; not
        ((logical2 arith2)
         (list *popr1* (list opcode ''r0 ''r1)))
        (arith1
         (list (list opcode ''r0 1)))
        (arith+-
         (if (and (eq f '-) (= (length args) 1)) (list '($neg 'r0 'r0)) ; unary minus
           (list *popr1* (list opcode ''r0 ''r1 ''r0))))
        (ash
         (let ((lab1 (gen-label)) (lab2 (gen-label)))
           (list '($neg 'r1 'r0) '($pop '(r0)) (list '$bpl lab1) '($neg 'r1 'r1) '($lsl 'r0 'r1) (list '$b lab2) lab1 '($asr 'r0 'r1) lab2)))
        (div
         (list *popr1* '($sdiv 'r0 'r1 'r0)))
        (mod
         (list *popr1* '($sdiv 'r12 'r1 'r0) '($mls 'r0 'r12 'r0 'r1)))
        (car
         (list '($ldr 'r0 '(r0 0)) '($ldr 'r0 '(r0 4))))
        (cdr
         (list '($ldr 'r0 '(r0 4))))
        (register
         (if (= (length args) 1) (list '($ldr 'r0 '(r0 0)))
           (list '($str 'r0 '(r1 0)))))
        (nth
         (let ((lab1 (gen-label)) (lab2 (gen-label)))
           (list *popr1* lab2 (list '$cbz ''r1 lab1) '($sub 'r1 1) '($ldr 'r0 '(r0 4)) (list '$b lab2) lab1 '($ldr 'r0 '(r0 0)) '($ldr 'r0 '(r0 4)))))
        (t
         (comp-function-call f args)))))))

(defun comp-function-call (f args)
  (append
   (when (> (length args) 1)
     (append 
      (list (reg-op '$mov (nth (1- (length args)) *params*) 'r0))
      (mappend
       #'(lambda (x) (list (list '$pop (list 'quote (list x)))))
       (reverse (subseq *params* 0 (1- (length args)))))))
   (list (list '$bl f))))

(defun comp-args (fn args k type env)
  (unless (or (null k) (= (length args) k))
    (error "Incorrect number of arguments to '~a'" fn))
  (let ((n (length args)))
    (mappend #'(lambda (y)
                 (let ((c (comp y env)))
                   (decf n)
                   (checktype fn type (code-type c))
                   (if (zerop n) (code c) (append (code c) '(($push '(r0)))))))
             args)))