;;; A multi-way tree for storing sparse indexed items. Although I just
;;; made this up, it looks like it is basically around in the
;;; literature, sometimes called a radix tree or a crit bit tree.  It
;;; is an alternative to a hash table or an array, and is, in some
;;; sense, a data structure that can be made more array-like or more
;;; tree-like. The index or key is always an integer, perhaps a
;;; bignum.

;;; The size of the nodes is parameterized: each internal node is a
;;; small array, say of length 8 =2^logsize Then <key, data> is
;;; stored in a node by decomposing the key into 3 =logsize bits at
;;; a time per layer in tree.  thus 30-bit keys will have depth 10
;;; max.  The initial sub-keys are the rightmost bits, and the remaining
;;; keys are computed by shifting the keys by logsize at each level.

;;; The key is never in fact stored in the tree, but computed by position.
;;; There is one subtlety worth noting. 

;;; Note that with logsize 3, we can store items with keys
;;; [0,1,...,7] in the topmost node.  If we need to store an item with
;;; key 9, it has the same trailing 3 bits as 1. To accomodate, we
;;; move the 1 down exactly one level.  That is, the topmost node
;;; looks like [0,p,...,7] where p=[1,9,nil,nil...nil].

;;; so an key like 1  be at the top level (or maybe 1 down).

;;; CON: Keys must be integers or mapped to integers.  Traversing the
;;; tree "in order" is tricky. Numerically adjacent keys will not be
;;; adjacent in the tree, except in the trivial case where all keys
;;; fit in one node.  Nodes may be largely empty.  A key with N bits
;;; will take about (N/logsize) probes.

;;; PRO: we need not know how long the longest key is, which would be
;;; the case if we were using a string-type decomposition starting
;;; from the left.  Short keys will be located near the top of the
;;; tree. The keys are not stored, saving space. By increasing logsize
;;; we can make the performance closer to that of an array.  We never
;;; compare keys; we only extract sub-keys and use them as indexes
;;; into arrays.

;;; author RJF June 13, 2008
;;; made faster, Oct 22, 2008, RJF
(eval-when (:compile-toplevel :load-toplevel :execute)
  (declaim (optimize (speed 3)(safety 1)))
 ;;   (defconstant logsize 5)		; intermediate node will have 2^logsize slots
  (defconstant logsize 3)		; best for debugging, not bad for running some tests.
)

(defmacro init-rtree()`(make-array #.(expt 2 logsize) :initial-element 0))
(defmacro leafp(x) `(numberp ,x))
(defmacro nodep(x) `(arrayp ,x))

(defun update-rtree(key val tree &aux node)
  ;; we use this for multiplication or addition of polynomials
  ;; but we could do anything with the update operation below
  ;; Declarations for speed.
  (declare (optimize (speed 3)(safety 0))
	   (type (simple-array t (#.(expt 2 logsize))) tree))
  
  (cond  ((fixnump key)(update-rtree-fix key val tree))
	 ;;above, optimized for fixnum key
	 ((< key #.(expt 2 logsize))
	  ;; we have found the level, or level-1 for the key.
	  (cond ((nodep (setf node (aref tree key)))
		 ;; if a node here, must descend one level further in tree
		 (update-rtree-fix 0 val node)); and update that node's location 0.
		(t(setf (aref tree key) (+ node val)))));; put it here.
	 
	 ;;The key has too many bits to insert at this level. 
	 ;;Compute the subkey and separate the rest of the key by masking and shifting.

	 (t (let ((ind(logand key #.(1- (expt 2 logsize)))))
	      ;; mask off the relevant bits for subkey
	      (declare (fixnum ind))
	      (setf node (aref tree ind))
	      (cond ((arrayp node)
		     (update-rtree (ash key #.(- logsize)) val node))
		    ;;descend into tree with rest of key.
		    (t (setf (aref tree ind)
			 (update-rtree 
			  (ash key #.(- logsize)) val 
			  (if (leafp node)
			      (update-rtree-fix 0 node (init-rtree))
			    (init-rtree)))))))))
  tree)

(defun query-rtree(key tree)		;works. though we don't use it..
  (declare (optimize (speed 3)(safety 0))
	   (type (simple-array t (#.(expt 2 logsize))) tree))
  (cond  ((< key #.(expt 2 logsize)) 
	  (if (arrayp (aref tree key)) (query-rtree 0 (aref tree key))
	    (aref tree key)))
	 (t (let* ((ind(logand key #.(1- (expt 2 logsize))))
		   (h (aref tree ind)))
	      (if h (query-rtree (ash key #.(- logsize)) h) nil)))))

(defun maptree (fn mt);; order appears jumbled, actually reversed-radix order
  ;; apply fn, a function of 2 arguments, to key and val, for each entry in tree.
  ;; order is "radix reversed"
  (declare (optimize (speed 3)(safety 0)))
  (labels 
      ((tt1 (path displace mt);; path = path to the key's location
	    (cond ((null mt) nil)
		  ((leafp mt) (funcall fn path mt));; function applied to key  and val.
		  (t;; must be an array
		   (do ((i 0 (1+ i)))
		       ((= i #.(expt 2 logsize)))
		     (declare (fixnum i))
		     (tt1 (+ (ash i displace) path) (+ displace logsize)(aref mt i)))))))
    (tt1 0 0  mt)))

#| example
(setf mt (init-rtree))
(defun t2()(setf mt (init-rtree-mt))(ir 1000 1000)(ir 2000 2000) (ir 1002 10002) (ordertree-nz mt))
(defun ir(a b)(insert-rtree a b mt))
(dotimes (i (expt 8 3))(update-rtree i i mt))
mt  ;;; a fully populated b-rtree with value n at location n, 0<= n < 8^3
(ordertree-nz mt)
|#

(defun mul-multivar-rtree(r s) 
  (declare(optimize (speed 3)(safety 0)))
  (multiple-value-bind (in out)(make-coders r s)
    (labels ((incode (z)(mapcar #'(lambda(h)(cons (car h)(funcall in (cdr h)))) z)))
        (setf r (ordertree-nz (mul-rtree (incode r) (incode s))))
	(dolist (h r r)(setf (car h)(funcall out (car h)))))))

(defun A2PA(A) ;; array of (exp . coef) to 2 arrays.
  (let* ((LA (length A))
	 (V nil)
	 (anse (make-array LA))
	 (ansc (make-array LA)))
    (dotimes (i LA (cons anse ansc))
      (setf V (aref A i))
      (setf (aref anse i) (car V))
      (setf (aref ansc i) (cdr V)))))
	    
    
  
(defun mul-rtree (r s)			;multiply two polynomials, in arrays
  (declare (optimize (speed 3)(safety 0)))
  (let* ((er (car r)) (es (car s))
	 (cr (cdr r)) (cs (cdr s))
	 (cri 0) (eri 0)
	 (ler (length er))
	 (les (length es))
	 (ans (init-rtree)) )
    (declare (type (simple-array integer (*)) cr cs er es)
	     (fixnum ler les))		;maybe some others are fixnums too..
    ;; do the NXM multiplies into the answer tree
    (dotimes (i ler (A2PA (ordertree-nz ans)))
      (declare (fixnum i))
      (setf cri (aref cr i) eri(aref er i))
      (dotimes (j les)
	(declare (fixnum j))
	(update-rtree (+ eri(aref es j)) (* cri(aref cs j)) ans)))))

(defun mul-rtree-nosort (r s);multiply two polynomials, in arrays, result in tree
  (declare (optimize (speed 3)(safety 0)))
  (let* ((er (car r)) (es (car s))
	 (cr (cdr r)) (cs (cdr s))
	 (cri 0) (eri 0)
	 (ler (length er))
	 (les (length es))
	 (ans (init-rtree)) )
    (declare (type (simple-array integer (*)) cr cs er es)
	     (fixnum ler les))		;maybe some others are fixnums too..
    ;; do the NXM multiplies into the answer tree
    (dotimes (i ler ans			;(L2PA (ordertree-nz ans))
		)
      (declare (fixnum i))
      (setf cri (aref cr i) eri(aref er i))
      (dotimes (j les)
	(declare (fixnum j))
	(update-rtree (+ eri(aref es j)) (* cri(aref cs j)) ans)
	))))

;;; The program update-rt1 below is a more general update function.
;;; First we show how we could use it for multiplication or addition.

(defun update-rtree+ (key val tree) ;; same as update-rtree, using rt1
  (update-rt1 key val tree #'(lambda(val oldval)(+ val (or oldval 0)))))

(defun update-rt1(key val tree fn)
  (labels((update-rtree
	   (key val tree)
	   ;; we use this for multiplication or addition of polynomials
	   ;; but we could do anything with the update operation below
	   ;; carefully, perhaps overly, declared for speed.
	   (declare (optimize (speed 3)(safety 0))
		    (type (simple-array t (#.(expt 2 logsize))) tree))
  
	   (cond  ((fixnump key)(update-rt1-fix key val tree fn))
		  ((< key #.(expt 2 logsize)) ; we have found the level, or level-1 for the key.
		   (cond ((nodep (aref tree (the fixnum key))); if a node here, must descend one level further in tree
			  (update-rtree 0 val (aref tree (the fixnum key)))); and update that node's location 0.
			 (t
			  ;; we have a leaf node to update.
			  (setf (aref tree (the fixnum key)) 
				(funcall fn val (aref tree (the fixnum key)))))))
		  ;;The key is still too big. Compute the subkey and the rest of the key by masking and shifting.
		  (t (let* ((ind(logand key #.(1- (expt 2 logsize)))); mask off the relevant bits for subkey
			    (h (aref tree ind)))
		       (declare (type (simple-array t (#.(expt 2 logsize))) tree)(fixnum ind))
		       (cond ((arrayp h)
			      (update-rtree (ash key #.(- logsize)) val h));descend into tree with rest of key.
			     ((leafp h)	; leaf, not link. We make a subtree and move the leaf down one level.
			      (setf (aref tree ind)
				    (update-rtree (ash key #.(- logsize)) val (update-rtree 0 h (init-rtree)))))
			     ;; h is nil means the node was empty. Just make a new subtree and insert.
			     (t (setf (aref tree ind)(update-rtree (ash key #.(- logsize)) val (init-rtree))))))))
	   tree))
    (update-rtree key val tree)))


(defun ordertree-nz(mt);;   order-tree the tree removing zeros.  sort after the whole tree is traversed
  (declare (optimize (speed 3)(safety 0))
	   (type (simple-array t (#.(expt 2 logsize))) tree))
  (let ((ans (make-array 10 :adjustable t :fill-pointer 0)))
     (declare   (type (array t (*)) ans))
  (labels 
      ((tt1 (path displace mt);; path = path to the key's location
	    (cond ((null mt) nil)
		  ;; line below changed to return nil if mt is zero
		  ((leafp mt)(unless (zerop mt) 
	       (vector-push-extend (cons path mt) ans)))
		  (t;; mt must be an array
		   (do ((i #.(1- (expt 2 logsize)) (1- i)))
		       ((< i 0) ans)
		      (declare (fixnum i) (type (array t (*)) mt))
		     (tt1 (+ (ash i displace) path) (+ displace logsize)(aref mt i)))))))
    (sort (tt1 0 0 mt) #'< :key #'car))))


(defun update-rtree-fix(key val tree &aux node)  ;;  optimization for  key being fixnum; minor hacks
  (declare (optimize (speed 3)(safety 0))
	   (type (simple-array t (#.(expt 2 logsize))) tree)
	   (fixnum key))
  
  (cond  ((< key #.(expt 2 logsize))	; we have found the level, or level-1 for the key.
	  (cond ((nodep (setf node (aref tree key))); if a node here, must descend one level further in tree
		 (update-rtree-fix 0 val node)); and update that node's location 0.
		(t(setf (aref tree key) (+ node val)))));; put it here.
	 
	  (t (let ((ind(logand key #.(1- (expt 2 logsize))))); mask off the relevant bits for subkey
	      (declare (fixnum ind))
	      (setf node (aref tree ind))
	      (cond ((arrayp node)
		     (update-rtree-fix (ash key #.(- logsize)) val node));descend into tree with rest of key.
		    (t (setf (aref tree ind)
			     (update-rtree-fix (ash key #.(- logsize)) val 
					   (if (leafp node)
					       (update-rtree-fix 0 node (init-rtree))
					     (init-rtree)))))))))  
  tree)

(defun update-rt1-fix(key val tree fn)
  (labels((update-rtree;;internally, use fixnum keys
	   (key val tree)
	   (declare (optimize (speed 3)(safety 0))
		    (type (simple-array t (#.(expt 2 logsize))) tree)
		    (fixnum key))
  
	   (cond  ((< key #.(expt 2 logsize)); we have found the level, or level-1 for the key.
		   (cond ((nodep (aref tree (the fixnum key))); if a node here, must descend one level further in tree
			  (update-rtree 0 val (aref tree (the fixnum key)))); and update that node's location 0.
			 (t
			  (setf (aref tree (the fixnum key)) 
				(funcall fn val (aref tree (the fixnum key)))))))
		  
		  
		  (t (let* ((ind(logand key #.(1- (expt 2 logsize)))); mask off the relevant bits for subkey
			    (h (aref tree ind)))
		       (declare (type (simple-array t (#.(expt 2 logsize))) tree)(fixnum ind))
		       (cond ((arrayp h)
			      (update-rtree (ash key #.(- logsize)) val h));descend into tree with rest of key.
			     ((leafp h)	; leaf, not link. We make a subtree and move the leaf down one level.
			      (setf (aref tree ind)
				    (update-rtree (ash key #.(- logsize)) val (update-rtree 0 h (init-rtree)))))
			     (t (setf (aref tree ind)(update-rtree (ash key #.(- logsize)) val (init-rtree))))))))
	   tree))
    (update-rtree key val tree)))