(in-package #:user-side-compiler)

;;; This is typically called a register allocator, but since the only real
;;; registers we're working with are special purpose, and there's some special
;;; 6502-related nonsense going on as well, it feels better to call it a value
;;; allocator instead.

(defparameter +accumulator-users+
  '(or ir-operation ir-assign ir-if))

(defun calls-exist-between-p (start end)
  (labels
      ((iter (now end)
         (cond
           ((eql now end)
            nil)
           ((typep now 'ir-call)
            now)
           ((null now)
            (error "Calls check crossed iblock boundary. This should not happen."))
           (t
            (iter (next now) end)))))
    (iter (next start) end)))

(defclass value-allocation ()
  ((%data :accessor data :initarg :data)
   (%strategy :accessor strategy :initform :temporary-variable)
   (%varvec-index :accessor varvec-index :initform nil)))

(defun optim-prepare-direct-instructions (iblock)
  (do-instructions (inst iblock)
    (when (typep inst 'ir-operation)
      (assert (= (length (inputs inst)) 2))
      (destructuring-bind (accumulator operand)
          (inputs inst)
        (when (typep (definition accumulator) '(or ir-fetchvar ir-getconst))
          (move-instruction-above (definition accumulator) inst))
        (when (typep (definition operand) '(or ir-fetchvar ir-getconst))
          (when (typep (definition operand) 'ir-fetchvar)
            (setf (inputs inst)
                  (substitute (input (definition operand)) operand
                              (inputs inst))))
          (delete-instruction (definition operand)))))))

(defmethod choose-allocation-strategy ((alc value-allocation))
  (setf
   (strategy alc)
   (let ((data (data alc)))
     (cond ((null (users data))
            :not-saved)
           ((typep data 'ir-variable)
            :named-variable)
           ((and (eql (next (definition data)) (last-use data))
                 (typep (last-use data) 'ir-if)
                 (typep (definition data) 'ir-comparison))
            :branch)
           ((and (eql (next (definition data)) (last-use data))
                 (typep (last-use data) +accumulator-users+)
                 (eql data (first (inputs (last-use data)))))
            :accumulator)
           ((and (= (length (users data)) 1)
                 (typep (last-use data) 'ir-call)
                 (not (calls-exist-between-p (definition data) (last-use data))))
            :direct-to-argvec)
           ((typep data 'ir-constant)
            :constant)
           (t
            :temporary-variable)))))

(defun allocate-values (start-iblock)
  (compute-lifetime-knowledge start-iblock)
  (let ((allocations '()))
    (do-iblocks (iblock start-iblock)
      (do-instructions (inst iblock)
        (unless (or (null (output inst))
                    (find (output inst) allocations :key #'data))
          (let ((allocation (make-instance 'value-allocation :data (output inst))))
            (choose-allocation-strategy allocation)
            (push allocation allocations)))
        (loop :for input :in (inputs inst)
              :unless (or (not (typep input 'ir-data)) (find input allocations :key #'data))
                :do (let ((allocation (make-instance 'value-allocation :data input)))
                      (choose-allocation-strategy allocation)
                      (push allocation allocations)))))
    (setf allocations (nreverse allocations))
    (let ((counter -1)
          (named (remove :named-variable allocations
                         :test-not #'eql :key #'strategy))
          (temporary (remove :temporary-variable allocations
                         :test-not #'eql :key #'strategy)))
      (loop :for allocation :in (append named temporary)
            :do (setf (varvec-index allocation) (incf counter))))
    allocations))

(defun optim-reuse-temporary-slots (start-iblock allocations)
  (let ((free '()))
    (do-iblocks (iblock start-iblock)
      (do-instructions (inst iblock)
        (let ((ending (find inst allocations
                            :key (alexandria:compose #'last-use
                                                     #'data)))
              (beginning (find (output inst) allocations :key #'data)))
          (cond ((and ending
                      (eql (strategy ending) :temporary-variable)
                      (not (null (varvec-index ending))))
                 (pushnew (varvec-index ending) free))
                ((and beginning
                      (eql (strategy beginning) :temporary-variable)
                      (not (null (varvec-index beginning)))
                      (not (null free)))
                 (setf (varvec-index beginning) (pop free)))))))))