#lang racket

(require data/bit-vector)

(struct rt-node [edge0 edge1 data] #:transparent #:mutable)
(struct rt-edge [label target] #:transparent)

(define (make-rt)
  (rt-node #f #f #f))

;; Helper functions for rt-node struct access based on whether the edge is 1 or 0
(define (rt-getter bit) (if bit rt-node-edge1 rt-node-edge0))
(define (rt-setter bit) (if bit set-rt-node-edge1! set-rt-node-edge0!))

;; Bit vector functions but short
(define bv-ref bit-vector-ref)
(define bv-len bit-vector-length)
(define bv-copy bit-vector-copy)
(define bv bit-vector)

(define (bv-common-len bv1 bv1-start bv2 bv2-start)
  ;; How to avoid this bv copy?? Why is there no drop for sequences either
  ;; Matthias pls
  (for/sum ([b1 (in-bit-vector (bv-copy bv1 bv1-start))]
            [b2 (in-bit-vector (bv-copy bv2 bv2-start))])
    #:break (not (equal? b1 b2))
    1))

(define (rt-partial-iterate node key [start 0])
  (cond
    [(>= start (bv-len key)) (list 'exact node)]
    [else
     (let* ([bit (bv-ref key start)]
            [getter (rt-getter bit)]
            [next-edge (getter node)]
            [next-label (and next-edge (rt-edge-label next-edge))]
            [next-target (and next-edge (rt-edge-target next-edge))]
            [next-common-len (and next-label (bv-common-len next-label 0 key start))])
       (cond
         [(and next-edge (= next-common-len (bv-len next-label)))
          (rt-partial-iterate next-target key (+ start next-common-len))]
         [next-edge (list 'partial node start next-edge next-common-len)]
         [else (list 'no-match node start)]))]))

(define (rt-insert! node key data)
  (define (insert-node! node key start data)
    (let* ([bit (bv-ref key start)]
           [setter! (rt-setter bit)])
      (setter! node (rt-edge (bv-copy key start) (rt-node #f #f data)))))
  (define (split-node! node key start orig-edge prefix-len data)
    (let* ([bit (bv-ref key start)]
           [setter! (rt-setter bit)]
           [orig-label (rt-edge-label orig-edge)]
           [orig-target (rt-edge-target orig-edge)]
           [new-orig-edge (rt-edge (bv-copy orig-label prefix-len) orig-target)]
           [new-insert-edge (rt-edge (bv-copy key (+ start prefix-len)) (rt-node #f #f data))]
           [diff-bit (bv-ref key (+ start prefix-len))]
           [common-node (rt-node (if diff-bit new-orig-edge new-insert-edge)
                                 (if diff-bit new-insert-edge new-orig-edge)
                                 #f)]
           [common-edge (rt-edge (bv-copy key start (+ start prefix-len)) common-node)])
      (setter! node common-edge)))
  (match (rt-partial-iterate node key)
    [(list 'exact node) (set-rt-node-data! data)]
    [(list 'partial node start orig-edge prefix-len)
     (split-node! node key start orig-edge prefix-len data)]
    [(list 'no-match node start)
     (insert-node! node key start data)]))

; (define test (make-rt))
; (define (test-insert! x)
;   (rt-insert! test (string->bit-vector x) x))
; (test-insert! "0001")
; (test-insert! "1000")
; (test-insert! "1010")
; (test-insert! "0011")
; (test-insert! "0000")
;
; (define (rt-dump node [prefix ""])
;   (displayln (format "~a node ~a" prefix (rt-node-data node)))
;   (define edge0 (rt-node-edge0 node))
;   (define edge1 (rt-node-edge1 node))
;   (when edge0
;     (displayln (format "~a  edge0 ~a" prefix (bit-vector->string (rt-edge-label edge0))))
;     (rt-dump (rt-edge-target edge0) (string-append prefix "  ")))
;   (when edge1
;     (displayln (format "~a  edge1 ~a" prefix (bit-vector->string (rt-edge-label edge1))))
;     (rt-dump (rt-edge-target edge1) (string-append prefix "  "))))
; (rt-dump test)