; Copyright (c) 1993-2007 by Richard Kelsey and Jonathan Rees. See file COPYING.

; Hilbert vectors are like vectors that grow as large as they need to.
; That is, they can be indexed by arbitrarily large nonnegative integers.

; The implementation allows for arbitrarily large gaps by arranging
; the entries in a tree.

; So-called because they live in an infinite-dimensional vector
; space...

; Modifications Sep 2022 Kon Lovett

;;

(include-relative "sparse-vectors-inlines")

;;

(: make-node-if-necessary (hilbert-node fixnum sparse-vector -> hilbert-node))
(: make-higher-if-necessary (sparse-vector integer -> hilbert-node))

(define (make-node-if-necessary node idx hilbert)
  (let ((val (node-ref node idx)))
    (if (node? val) val
        (let ((new (make-node hilbert)))
          (node-set! node idx new)
          (set-hilbert-capacity! hilbert (fx+ (hilbert-capacity hilbert) 1))
          new ) ) ) )

(define (make-higher-if-necessary hilbert index)
  ;in 1st node?
  (if (< index (hilbert-node-size hilbert))
      ;then not necessary
      (hilbert-root hilbert)
      ;else push everything under
      (let ((new-root (make-node hilbert)))
        (node-set! new-root 0 (hilbert-root hilbert))
        (set-hilbert-capacity! hilbert (fx+ (hilbert-capacity hilbert) 1))
        (set-hilbert-root! hilbert new-root)
        (set-hilbert-height! hilbert (fx+ (hilbert-height hilbert) 1))
        (let* ((right-bits (fxneg (hilbert-bits hilbert)))
               (index (node-size-expt index right-bits)) )
          (make-node-if-necessary (make-higher-if-necessary hilbert index)
                                  (bitwise-and index (hilbert-mask hilbert))
                                  hilbert) ) ) ) )

;;

(define (*sparse-vector-get hilbert index handler)
  (let ((right-bits (fxneg (hilbert-bits hilbert)))
        (mask (hilbert-mask hilbert)) )
    (let recur ((height (hilbert-height hilbert)) (index index))
      (if (fx= height 1) (handler hilbert (hilbert-root hilbert) index)
          (let ((node (recur (fx- height 1) (node-size-expt index right-bits))))
            (handler hilbert node (bitwise-and index mask))) ) ) ) )

(define (node-ref* hilbert node idx)
  (if (node-ref-valid? node idx) (node-ref node idx)
      (hilbert-default hilbert)) )

;; Public (Unchecked)

(define (*sparse-vector-ref hilbert idx)
  (hilbert-def-value* (*sparse-vector-get hilbert idx node-ref*)) )

(define (*sparse-vector-set! hilbert index value)
  (let* ((right-bits (fxneg (hilbert-bits hilbert)))
         (mask (hilbert-mask hilbert))
         (node (let recur ((height (hilbert-height hilbert)) (index index) )
                 (if (fx= height 1) (make-higher-if-necessary hilbert index)
                     (let ((index (node-size-expt index right-bits)))
                       (make-node-if-necessary
                        (recur (fx- height 1) index)
                        (bitwise-and index mask)
                        hilbert) ) ) ) )
         (idx (bitwise-and index mask) ) )
    ;maintain population count
    (when (hilbert-def? (node-ref node idx))
      (set-hilbert-count! hilbert (fx+ (hilbert-count hilbert) 1)) )
    (node-set! node idx value) ) )

(define (*sparse-vector-unset! hilbert index silent?)
  (let ((last-node (the (or false hilbert-node) #f))
        (last-idx (the (or false fixnum) #f)) )
    (define (last-node-ref hilbert node idx)
      (if (node-ref-valid? node idx)
          (begin
            (set! last-node node)
            (set! last-idx idx)
            (node-ref node idx) )
          (hilbert-default hilbert)) )
    (let ((x (*sparse-vector-get hilbert index last-node-ref)) )
      (if (not (hilbert-def? x))
          ;then unbind
          (begin
            #;(assert (node? last-node))
            #;(assert (fixnum? last-idx))
            (node-unset! last-node last-idx (hilbert-default hilbert))
            ;maintain population count
            (set-hilbert-count! hilbert (fx- (hilbert-count hilbert) 1)) )
          ;else a problem?
          (unless silent?
            (error-bounds 'sparse-vector-unset! hilbert index)) ) ) ) )

;; Public

(define (make-sparse-vector . rest)
  (let-optionals rest ((default #f) (bits DEFAULT-HILBERT-BITS))
    (let ((hilbert (make-hilbert 1 #f (make-hilbert-def default) 0 0 bits)))
      (set-hilbert-root! hilbert (make-node hilbert))
      (set-hilbert-capacity! hilbert (fx+ (hilbert-capacity hilbert) 1))
      hilbert ) ) )

(define (sparse-vector? x) (hilbert? x))

(define check-sparse-vector (cut check-structure <> <> sparse-vector))

(define (sparse-vector-count hilbert)
  (hilbert-count (check-sparse-vector 'sparse-vector-count hilbert)) )

(define (sparse-vector-set! hilbert index value)
  (*sparse-vector-set! (check-sparse-vector 'sparse-vector-set! hilbert)
                       (check-exact-uinteger 'sparse-vector-set! index)
                       (check-bound-value 'sparse-vector-set! value 'value)) )

(define sparse-vector-ref
  (getter-with-setter
   (lambda (hilbert index)
     (*sparse-vector-ref
      (check-sparse-vector 'sparse-vector-ref hilbert)
      (check-exact-uinteger 'sparse-vector-ref index)))
   sparse-vector-set!))

;FIXME unset does not scavenge unneeded nodes

(define (sparse-vector-unset! hilbert index . rest)
  (*sparse-vector-unset! (check-sparse-vector 'sparse-vector-unset! hilbert)
                         (check-exact-uinteger 'sparse-vector-unset! index)
                         (optional rest #t)) )