;;; -*- Mode:Common-Lisp; Package:L; Base:10 -*-
;;;
;;; an implementation of register coloring
;;;
;;; ayers 5/7-8/88
;;;
;;; Notes:
;;;  -- This code is based upon register coloring ideas from
;;;     Chaitin's paper "Register allocation via coloring"
;;;     [ Register Allocation via Coloring, G. J. Chaitin et. al., ACM
;;;       Transactions on Computer Languages 6 (1981), pp. 47-57. ]
;;;  -- I have not put much effort into efficiency issues
;;;  -- Mechanism for handling LOCKS, S-STORES, and S-LOADS is
;;;     mostly absent. There are some complications and implications
;;;     that I will probably turn into an L memo, dealing with
;;;     STATIC vars, etc.
;;;  -- The driver is not very intelligent, and not especially
;;;     efficient. It could re-use a lot of information from
;;;     earlier attempts, and could avoid trying colorings that
;;;     will obviously fail (i.e. 8-colorings of things with
;;;     static variables)
;;;
;;; Converted for il-node use.
;;; jsp 20-May-88



;;; debugging macro

(defvar *color-debug* nil)

(defmacro when-debug (form)
  `(if *color-debug* ,form))

;;;; Registers
;;;
;;; a register data structure
;;;
;;; <interferences> is a list of other registers that the current register
;;;   cannot be coalesced with
;;; <interference-count> is the length of the interference list
;;; <real> is the register # if this is  a real register (i.e. %R3 has a 3 here)
;;;   otherwise <real> is nil.
;;; <synonyms> is a list of other registers that have been merged into this one
;;;   through allocation actions
;;; <alias> is a register that this has been merged into, i.e. a forwarding pointer.
;;; <uses> is a count of how many times this register appears in the code. it is used
;;;   to discount registers that don't actually need space.

(defstruct (register (:print-function print-register))
  (name 'unnamed :type symbol)
  (interferences '() :type list)
  (interference-count 0 :type integer)
  (uses 0 :type integer)
  (synonyms '() :type list)
  (alias '())
  (real '()))

(defun print-register (x s d)
  (ignore d)
  (format s "{~s}" (register-name x)))

;;;; Graphs
;;;
;;; interference graph data structure
;;;
;;; <name> is used for printing purposes only. the code below sets the name
;;;   to be the name of the function undergoing allocation.
;;; <all-registers> holds a list of every register
;;;   that has been created for this graph. these are remembered primarily
;;;   to keep the aliasing information available.
;;; <registers> holds the active registers. these are the registers that 
;;;   are still under consideration in the allocation process.
;;; <symbolic-registers> holds all the (used) symbolic registers for the
;;;   current code. this field is used to allow allocation to backtrack quickly
;;;   the event that initial coloring efforts fail.

(defstruct (graph (:print-function print-graph))
  (name 'unnamed :type symbol)
  (all-registers '() :type list)
  (symbolic-registers '() :type list)
  (registers '() :type list))

(defun print-graph (graph s d)
  (ignore d)
  (format s "{graph ~s}" (graph-name graph)))

;;;; Graph/Register operations
;;;
;;; add a new interference.
;;;
;;; used, for example, when we discover that reg1 and reg2 must be alive
;;; simultaneously.

(defun add-interference (reg1 reg2)
  "add interference between reg1 and reg2"
  (unless (member reg1 (register-interferences reg2))
    (push reg1 (register-interferences reg2))
    (incf (register-interference-count reg2)))
  (unless (member reg2 (register-interferences reg1))
    (push reg2 (register-interferences reg1))
    (incf (register-interference-count reg1))))

;;; merge reg1 into reg2.
;;;
;;; used to cause reg1 to be allocated into the same register as
;;; reg2.
;;;
;;; reg1's interferences fall into two classes: those shared by
;;; reg2 and those not shared. For class 1, we simply delete reg1 from
;;; each reg and reduce the count by 1. For class 2, we substitute reg2
;;; for reg1 and leave the count unchanged.
;;;
;;; the set-difference/intersection operations can probably be done more
;;; efficiently

(defun coalesce (reg1 reg2 graph)
  "merge reg1 into reg2, in graph. Returns graph"
  (if (eq reg1 reg2)
      (when-debug (format t "~&Trivial coalesce: ~s" reg1))
      (nontrivial-coalesce reg1 reg2 graph)))

(defun nontrivial-coalesce (reg1 reg2 graph)
  (let ((r1int (register-interferences reg1))
	(r2int (register-interferences reg2)))
    (let ((shared-interferences (intersection r1int r2int :test #'eq))
	  (added-interferences (set-difference r1int r2int :test #'eq)))
      (dolist (i added-interferences)
	(update-interference i reg1 reg2))
      (dolist (i shared-interferences)
	(remove-interference i reg1))
      (setf (register-interferences reg2)
	    (append added-interferences (register-interferences reg2)))
      (incf (register-interference-count reg2) (length added-interferences))))
  (push reg1 (register-synonyms reg2))
  (setf (register-alias reg1) reg2)
  (setf (graph-registers graph) (delete reg1 (graph-registers graph) :test #'eq))
  graph)

(defun update-interference (in-reg old new)
  "substitute new for old in the interference list of in-reg"
  (setf (register-interferences in-reg)
	(substitute new old (register-interferences in-reg) :test #'eq)))

(defun remove-interference (in-reg reg)
  "remove reg from the interference list of in-reg"
  (setf (register-interferences in-reg)
	(delete reg (register-interferences in-reg) :test #'eq))
  (decf (register-interference-count in-reg)))

;;; predicate to check for interference

(defun interference? (reg1 reg2)
  (member reg2 (register-interferences reg1)))

;;;; find symbolic registers in a function
;;;
;;; given a function header, discover all of the
;;; symbolic registers inside that may need allocation
;;; (this set may actually include some registers that are
;;; not used by the code)

(defun produce-symbolic-registers (fn)
  (mapcar #'(lisp:lambda (x) (make-register :name (declaration-symbol x)))
	  (extract-symbolic-registers fn)))

(defun extract-symbolic-registers (fn)
  (let ((sr)
;	(dl (list KIND-ARG KIND-LOCAL KIND-TEMP KIND-CONST)))
	(dl (list KIND-ARG KIND-LOCAL KIND-TEMP)))
    (maphash #'(lisp:lambda (ignore decl)
		 (when (member (declaration-kind decl) dl)
		   (push decl sr)))
	     (attribute-get (fn-et fn) :declarations))
    sr))

;;;; use counts
;;;
;;; given the set of registers extracted from a function,
;;; and the code for that function, 
;;; discover which registers are actually USED by the code.
;;; Side effects into the USE field of the registers in the register-list

(defun record-use-count (code register-list)
  (mapc #'(lisp:lambda (i) (record-uses i register-list)) code))

;;; what are the symbolic registers mentioned in this instruction?

(defun find-symregs-in (inst)
  (remove-if-not #'(lisp:lambda (sym)
		     (and (declaration? sym)
			  (neq (declaration-kind sym) KIND-CONST)))
		 (append (il-node-read inst) (il-node-written inst))))

;;; go from declaration to register. this should never fail.

(defun sym->reg (sym reg-set)
  (or (find-if #'(lisp:lambda (x) (eq (register-name x) sym)) reg-set)
      (error "~&Unknown register ~s" sym)))

(defun decl->reg (decl reg-set)
  (sym->reg (declaration-symbol decl) reg-set))

;;; record uses in an instruction.

(defun record-uses (inst reg-list)
  (mapcar #'(lisp:lambda (x) (incf (register-uses x)))
	  (mapcar #'(lisp:lambda (x) (decl->reg x reg-list)) (find-symregs-in inst))))

;;; this function removes use-less registers from
;;; a list of registers.

(defun remove-useless-registers (reg-list)
  (remove-if #'(lisp:lambda (x) (zerop (register-uses x)))
	     reg-list))

;;;; real register interference graph
;;;
;;; to enforce the convention that the real registers are
;;; distinct, we create a register graph, forcing each register to
;;; interfere with all the others. The parameter <n> controls how
;;; many registers we wish to consider the machine to have.
;;;
;;; this set could probably be built more efficiently.

(defun build-register-interference-graph (n)
  (let ((registers))
    (dotimes (i n)
      (push (make-register :name (intern (format nil "%R~d" i)) :real i) registers))
    (dolist (r registers)
      (setf (register-interferences r) (set-difference registers (list r)))
      (setf (register-interference-count r) (- n 1)))
    registers))

;;;; graph creation
;;;
;;; build an interference graph, given the list of registers that
;;; appear, and a trial size for our coloring.

(defun build-interference-graph (fn n)
  (let ((g (make-graph :name (fn-name fn)))
	(s (produce-symbolic-registers fn))
	(r (build-register-interference-graph n))
	(code (fn-il-nodes fn)))
    (record-use-count code s)
    (setf (graph-symbolic-registers g) s
	  (graph-all-registers g) (append s r)
	  (graph-registers g) (append (remove-useless-registers s) r))
    g))

;;;; handle static information
;;;
;;; this routine takes any STATIC declarations in the 
;;; function into account. [recall that these declarations 
;;; indicate that the mentioned variables must have the same
;;; allocation for the extent of the function].
;;;
;;; the strategy used is to make STATIC variables interfere with all
;;; other symbolic registers, thereby assuring that their allocation
;;; will have the proper extent.  Static variables must have at least
;;; one use, so this processing can be done after use analysis; since
;;; the static interferences will be a superset of any normal ones, this
;;; analysis should also take place after taking instruction
;;; interferences into account.
;;;
;;; this strategy may seem somewhat aggresive, because a static local
;;; does not need to lock down a slot until it's value is calculated.
;;; we could amend this to allow static locals to merely interfere with
;;; everthing live after their definition
;;; point, but this information seems hard to come by.
;;;
;;; because %R0-%R7 and %R24-%R31 are used specially, static variables
;;; cannot live in any of these, and so these interferences are added as
;;; well.  this formulation of the static rule allows us to use
;;; %R24-%R31 as temporaries (not currently done, because the LOCK/SLOAD
;;; rules are not in place yet.) An alternative is to let statics live
;;; in %R0-%R7, and never mess with %R24-%R31, but this
;;; seems more restrictive.  Some experimentation would probably pay off here.

(defun add-static-information (graph fn)
  (let ((regs (graph-registers graph)))
    (let ((rregs (remove-if-not '%R0-%R7-or-%R24-%R31 regs))
	  (sregs (remove-if 'register-real regs)))
      (let ((sl (mapcar #'(lisp:lambda (x) (sym->reg x sregs)) (get-static-list fn))))
	(when-debug (format t "~&Static: ~{~s ~}" sl))
	(when-debug (format t "~&Reals: ~{~s ~}" rregs))
	(mapcar #'(lisp:lambda (x) (staticize x sregs rregs)) sl)))))

(defun %R0-%R7-or-%R24-%R31 (reg)
  (let ((x (register-real reg)))
    (and x
	 (or (<= 0 x 7)
	     (<= 24 x 31)))))

;;; can probably be made more efficient

(defun staticize (sreg s-registers r-registers)
  (let ((all-but-me (delete sreg s-registers :test #'eq)))
    (let ((ints (append all-but-me r-registers)))
      (mapcar #'(lisp:lambda (x) (add-interference x sreg)) ints))))

;;; assumes only 1 static list exists.
;;; (plows through lots of list structure to get to
;;; intermediate-language instructions)

(defun get-static-list (fn)
  (let ((sl (find-if #'(lisp:lambda (x) (eq (car x) 'static))
		     (cadr (cdar (cdar (fn-normal-entry-vector fn)))))))
    (if sl
	(cadr sl)
	(error "~&No STATIC declaration in ~s" fn))))

;;;; add interference information
;;;
;;; take any interferences caused by instructions into account
;;; here instructions are IL-NODEs.
;;; information is recorded by side effecting the structures in
;;; graph.
;;;
;;; for each il-node, we extract the live list. we then call add-conflicts
;;; which adds an interference between each pair of registers on the live list.

(defun add-interference-information (graph code)
  (when-debug (format t "~%Extracting liveness information ..."))
  (mapc #'(lisp:lambda (x) (add-instruction-interferences x graph)) code))

(defun add-instruction-interferences (instruction graph)
  (let ((regs (graph-all-registers graph)))
    (let ((liveregs (mapcar #'(lisp:lambda (x) (decl->reg x regs))
			    (set-difference (il-node-alive instruction) (il-node-newly-dead instruction))
			    )))
      (when-debug (format t "~& ~3d ~s ++ ~{ ~s ~}" instruction (il-node-n instruction) liveregs))
      (maplist #'add-conflicts liveregs)
      (add-any-other-interferences instruction regs liveregs))))

(defun add-conflicts (decl-list)
  (let ((a (first decl-list)))
    (dolist (b (rest decl-list))
      (add-interference a b))))

;;; this function handles any special interferences on a per-instruction
;;; basis.
;;;
;;; LEAP: leap is used to handle tr calls to other functions. when the leap is
;;; compiled, a special value is set in the "leap destination" declaration
;;; to indicate the proper interferences that this register must have to prevent
;;; errors in allocation.
;;;
;;; MOVE: when moving a symbolic register to a real one, look for the
;;; interference count hack attribute, and add interferences if necessary.
;;; these moves arise in tail recursive calls and in function returns. only
;;; the first kinds have this special attribute.

(defun add-any-other-interferences (il-node regs liveregs)
  (ignore liveregs)
  (select (il-node-op il-node)
    (*il-goto-instruction*
     (if (neq (declaration-kind (first (il-node-targets il-node)))
	      KIND-LABEL)
	 (add-interference-to-first-n-registers (first (il-node-targets il-node)) regs)))
    (*il-move-instruction*
     (mapc #'(lisp:lambda (reg)
	       (add-interference-to-first-n-registers reg regs))
	   (il-node-read il-node)))))

(defun add-interference-to-first-n-registers (decl regs)
  (let ((count (attribute-get decl :interference-count-hack)))
    (when (numberp count)
      (let ((real-regs '(%R0 %R1 %R2 %R3 %R4 %R5 %R6 %R7))
	    (reg (decl->reg decl regs)))
	(when-debug (format t "~&Adding ~d primitive interference~:p to ~s" count reg))
	(dotimes (i count)
	  (add-interference reg (sym->reg (pop real-regs) regs)))))))

;;;; coalescing phase
;;;
;;; after we have built the graph, and recorded all of the
;;; interference information present due to code constraints.
;;; try to do some preliminary merging of registers.
;;;
;;; specifically, if we see [move a b] then try to get a and b
;;; into the same register. other optimizations of this kind are
;;; possible, but not implemented.

(defun handle-coalesces (graph instructions)
  (dolist (i instructions)
    (select (il-node-op i)
      (*il-move-instruction*   (handle-move   i graph))
      (*il-argdef-instruction* (handle-argdef i graph))
      (*il-return-instruction* (handle-return i graph))))
  graph)



;;; handle a move instruction.

(defun handle-move (il-node graph)
  (let ((source (first (il-node-read il-node)))
	(target (first (il-node-written il-node))))

  (when-debug (format t "~&Looking at ~s" il-node))
  (if (symbolp source)
      (let ((real-reg (d-or-s->r source graph))
	    (sym-reg (d-or-s->r target graph)))
	(when (and real-reg (not (interference? real-reg sym-reg)))
	  (when-debug (format t "~&Merging ~s into ~s" sym-reg real-reg))
	  (coalesce sym-reg real-reg graph)))
      (if (symbolp target)
	  (let ((real-reg (d-or-s->r target graph))
		(sym-reg (d-or-s->r source graph)))
	    (when (and real-reg (not (interference? real-reg sym-reg)))
	      (when-debug (format t "~&Merging ~s into ~s" sym-reg real-reg))
	      (coalesce sym-reg real-reg graph)))
	  (if (neq (declaration-kind source) KIND-CONST)
	      (let ((rr1 (d-or-s->r source graph))
		    (rr2 (d-or-s->r target graph)))
		(when (not (interference? rr1 rr2))
		  (if (register-real rr2)
		      (progn
			(when-debug (format t "~&Merging ~s into ~s" rr1 rr2))
			(coalesce rr1 rr2 graph))
		      (progn
			(when-debug (format t "~&Merging ~s into ~s" rr2 rr1))
			(coalesce rr2 rr1 graph))))))))))



;;; handle an argdef instruction.

(defun handle-argdef (il-node graph)
  (let* ((n 0)
	 (sources)
	 (targets (il-node-defined il-node)))
    
    (dolist (target targets)
      (push (make-register :name (intern (format nil "%R~d" n)) :real n) sources)
      (incf n))
    (setq sources (nreverse sources))
    
    (do ((source (pop sources) (pop sources))
	 (target (pop targets) (pop targets)))
	((null source))
      
      ;; We know that the source is a REGISTER, because we've just created it, and much of the checking in
      ;; handle-move becomes superfluous.
      
      (when-debug (format t "~&Looking at ~s" il-node))
      (if (symbolp target)
	  (let ((real-reg (d-or-s->r target graph))
		(sym-reg source))
	    (when (and real-reg (not (interference? real-reg sym-reg)))
	      (when-debug (format t "~&Merging ~s into ~s" sym-reg real-reg))
	      (coalesce sym-reg real-reg graph)))
	  (let ((rr1 source)
		(rr2 (d-or-s->r target graph)))
	    (when (not (interference? rr1 rr2))
	      (if (register-real rr2)
		  (progn
		    (when-debug (format t "~&Merging ~s into ~s" rr1 rr2))
		    (coalesce rr1 rr2 graph))
		  (progn
		    (when-debug (format t "~&Merging ~s into ~s" rr2 rr1))
		    (coalesce rr2 rr1 graph)))))))))



;;; handle a return instruction.

(defun handle-return (il-node graph)
  (let* ((n 0)
	 (sources (il-node-read il-node))
	 (targets))
    
    (dolist (source sources)
      (push (make-register :name (intern (format nil "%R~d" n)) :real n) targets)
      (incf n))
    (setq targets (nreverse targets))
    
    (do ((source (pop sources) (pop sources))
	 (target (pop targets) (pop targets)))
	((null source))
      
      ;; We know that the target is a REGISTER, because we've just created it, and much of the checking in
      ;; handle-move becomes superfluous.
      
      (when-debug (format t "~&Looking at ~s" il-node))
      (if (symbolp source)
	  (let ((real-reg (d-or-s->r source graph))
		(sym-reg target))
	    (when (and real-reg (not (interference? real-reg sym-reg)))
	      (when-debug (format t "~&Merging ~s into ~s" sym-reg real-reg))
	      (coalesce sym-reg real-reg graph)))
	  (if (neq (declaration-kind source) KIND-CONST)
	      (let ((rr1 (d-or-s->r source graph))
		    (rr2 target))
		(when (not (interference? rr1 rr2))
		  (if (register-real rr2)
		      (progn
			(when-debug (format t "~&Merging ~s into ~s" rr1 rr2))
			(coalesce rr1 rr2 graph))
		      (progn
			(when-debug (format t "~&Merging ~s into ~s" rr2 rr1))
			(coalesce rr2 rr1 graph))))))))))




;;;; go from a symbol to register, in a graph
;;; 
;;; these routines map a symbol or declaration into an ACTIVE graph register.
;;;
;;; method: look for a register with the same name. if we find one,
;;; chase the alias list until unaliased.
;;;
;;; note: "undefined" things pass through, returning NIL. This represents a minor
;;; inefficiency (having to do with %R24, etc.) that will be fixed up in the future.
;;; at that point, undefined things will cause errors.

(defun find-register (sym reg-list)
  (let ((original (find-if #'(lisp:lambda (reg) (eq sym (register-name reg))) reg-list)))
    (when original
      (de-alias original))))

(defun de-alias (reg)
  (if (register-alias reg) (de-alias (register-alias reg)) reg))

(defun sym->register (sym graph)
  (or (find-register sym (graph-all-registers graph))
      (progn
	(when-debug (format t "~&Can't find register ~s" sym))
	'())))

(defun decl->register (decl graph)
  (sym->register (declaration-symbol decl) graph))

(defun d-or-s->r (x graph)
  (typecase x
    (symbol (sym->register x graph))
    (declaration (decl->register x graph))))

;;;; (try to) color the graph
;;;
;;; now, given the coalesced graph, can we color it in N colors?
;;; returns graph with only real registers.
;;; 
;;; approach:
;;;  split graph into real and symbolic registers (ranked by #of conflicts).
;;;  given a symbolic register, try to merge it in succession with real registers, until
;;;  it works. repeat until there are no un-merged symbolic registers.

(defun trivial-color (graph)
  (let ((regs (graph-registers graph)))
    (let ((symregs (sort (remove-if #'register-real regs)
		       #'(lisp:lambda (x y) (<= (register-interference-count x)
						(register-interference-count y)))))
	  (realregs (remove-if-not #'register-real regs)))
      (dolist (sym symregs)
	(unless (dolist (real realregs)
		  (when (not (interference? real sym))
		    (coalesce sym real graph)
		    (return t)))
	  (progn
	    (when-debug (format t "~&Failed to color ~s" sym))
	    (throw 'color-failed '()))))
      graph)))

;;;; graph detailed printing

(defun graph->map (graph)
  (mapcar 'reg->maplist (graph-registers graph)))

(defun reg->maplist (reg)
  (list* (register-real reg) (mapcar 'register-name (register-synonyms reg))))

(defun describe-graph (graph)
  (format t "~%INTERFERENCE GRAPH")
  (format t "~%==================")
  (format t "~%~%active registers:")
  (mapc #'print-register-entry (graph-registers graph))
  graph)

(defun print-register-entry (r)
  (format t "~& ::: register ~s (~d)" (register-name r) (register-interference-count r))
  (format t "~&~{~< ~1,80:; ~s~>~^ ~}" (register-synonyms r)))

;;;; driver function (testing)
;;;
;;; pass in a symbol, which is a function name.
;;;
;;; notes: does not use any information present in an old graph to help
;;; retries, even though some of the information in a graph is put there
;;; explicitly for this purpose (an efficiency hack).

(defun coloring-test (sym)
  (let ((fn (get-value sym)))
    (let ((code (fn-il-nodes fn)))
      (describe-graph
	(do ((count 8 (+ count 8)))
	    ((= count 32))
	  (when-debug (format t "~&Trying an ~r-coloring ... " count))
	  (let ((attempt (catch 'color-failed
			   (let ((graph (build-interference-graph fn count)))
			     (add-interference-information graph code)
			     (add-static-information graph fn)
			     (handle-coalesces graph code)
			     (trivial-color graph)))))
	    (when attempt (return attempt))))))))

;;;; allocate registers via coloring

(defun allocate-registers-via-coloring (fn)
  (let ((color-graph (colorize fn)))
    (if color-graph
	(coloring-finalize color-graph fn)
	(error "~&Register coloring failed..."))))

(defun colorize (fn)
  (let ((code (fn-il-nodes fn)))
    (do ((count 2 (+ count 2)))
	((= count 32))
;    (do ((count 8 (+ count 8)))
;	((= count 32))
      (when-debug (format t "~&Trying an ~r-coloring ... " count))
      (let ((attempt (catch 'color-failed
		       (let ((graph (build-interference-graph fn count)))
			 (add-interference-information graph code)
			 (add-static-information graph fn)
			 (handle-coalesces graph code)
			 (trivial-color graph)))))
	(when attempt (return attempt))))))


;;;; finish up (interface to backend)
;;;
;;; write allocation information into each of the symbolic registers.
;;; output preamble appropriate for our register usage.

(defun coloring-finalize (graph fn)
  (map-allocations-to-declarations graph fn)
  (add-preamble graph fn))

(defun map-allocations-to-declarations (graph fn)
  (mapc #'(lisp:lambda (x) (map-reg-to-decl x fn graph)) (graph-symbolic-registers graph)))

(defun map-reg-to-decl (reg fn graph)
  (when (plusp (register-uses reg))
    (let ((symbol (register-name reg)))
      (let ((real-identity (d-or-s->r symbol graph)))
	(let ((real-index (register-real real-identity)))
	  (let ((decl (get-local-declaration symbol fn)))
	    (setf (declaration-allocation decl) real-index)))))))

(defun add-preamble (graph fun-value)
  (let ((n-regs (length (graph-registers graph)))
;	(fun-code (fn-il-nodes fun-value))
	)
    (when-debug (format t "~& >> Colored ~s in ~d colors" graph n-regs))
    (when-debug (describe-graph graph))
    (setf (fn-n-temporaries fun-value) n-regs)
    ;; add 1 arg bank for 16 regs, 2 for 24.
;    (dotimes (i (1- (truncate n-regs 8)))
;      (let ((first-node (car fun-code))
;	    (new-node (new-il-node :intermed (list 's-alloc-i %#data (+ i 5)))))
;	(add-offspring new-node first-node)
;	(add-ancestor  first-node new-node)
;	(push new-node fun-code)))
;    (setf (fn-il-nodes fun-value) fun-code)
    ))