(define-class <pvector> ()
  ((size :init-keyword :size)
   (shift :init-keyword :shift)
   (root :init-keyword :root)))

(define pv:null (make <pvector> :size 0 :shift 0 :root #()))

(define (update-vector vec i val)
  (let ((tr (vector-copy vec)))
    (vector-set! tr i val)
    tr))

(define (expand-vector vec val)
  (let* ((length (vector-length vec))
	 (tr (vector-copy vec 0 (+ 1 length) #f)))
    (vector-set! tr length val)
    tr))


(define (pvector . args)
  (fold pv:cons pv:null args))


(define (pv:assoc pvec i val)
  (let ((shift (slot-ref pvec 'shift)))
      (make <pvector>
	:size (slot-ref pvec 'size)
	:shift shift
	:root (pv:assoc* shift
			 (slot-ref pvec 'root)
			 i
			 val))))


(define (pv:assoc* shift vec i val)
  (if (zero? shift)
      (update-vector vec (logand i #x01f) val)
      (let* ((subidx (logand (ash i (- shift)) #x01f))
	     (child (pv:assoc* (- shift 5) (vector-ref vec subidx) i val)))
	(update-vector vec subidx child))))


(define (pv:nth pvec i)
  (let loop ((shift (slot-ref pvec 'shift))
	     (arr   (slot-ref pvec 'root)))
    (if (zero? shift)
	(vector-ref arr (logand i #x01f))
	(loop (- shift 5) (vector-ref arr (logand (ash i (- shift)) #x01f))))))


(define (pv:cons obj pvec)
  (let* ((shift (slot-ref pvec 'shift))
	 (size  (slot-ref pvec 'size))
	 (root  (slot-ref pvec 'root))
	 (new-root (pv:cons* shift root obj)))
    (if new-root
	(make <pvector>
	  :size (+ size 1)
	  :shift shift
	  :root new-root)
	(make <pvector>
	  :size (+ size 1)
	  :shift (+ shift 5)
	  :root (vector root (make-vector-chain (+ shift 5) obj))))))

(define (pv:cons* shift vec val)
  (if (zero? shift)
      (pv:cons-leaf vec val)
      (pv:cons-node shift vec val)))


(define (node-full? vec)
  (= (vector-length vec) 32))

(define (pv:cons-leaf vec val)
  (if (node-full? vec)
      #f
      (expand-vector vec val)))

(define (pv:cons-node shift vec val)
  (let ((child (pv:cons* (- shift 5)
			 (vector-ref vec (- (vector-length vec) 1))
			 val)))
    (if child
	(update-vector vec (- (vector-length vec) 1) child)
	(if (node-full? vec)
	    #f
	    (expand-vector vec (make-vector-chain shift val))))))

(define (make-vector-chain shift val)
  (if (= shift 5)
      (vector val)
      (vector (make-vector-chain (- shift 5) val))))