#lang racket

(require "iputil.rkt")

(provide make-rt rt-add! rt-del! rt-lookup rt-dump rt-flatten)

;; Radix tree implementation for efficient routing lookups
;; This module provides a radix tree type indexed by bit-lists (lists of booleans representing IP
;; address bits), which supports arbitary node data at any tree depth. Lookup finds all the matches.

;; An RT is a (rt-node RTedge RTedge [Listof Any])
;; An RTkey is a [Listof Bool]
;; An RTedge is a (rt-edge RTkey RT)
(struct rt-node [edge0 edge1 data] #:transparent #:mutable)
;; Edge labels always contain all of their bits, even though the first bit is redundant
;; (represented by whether this was an edge0 or edge1 in the parent node)
(struct rt-edge [label target] #:transparent)

;; -> RT
;; Creates a new empty RT
(define (make-rt)
  (rt-node #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!))

;; RTkey RTkey -> Nat
;; Computes the length of the longest common prefix in bl1 and bl2
(define (bl-common-len bl1 bl2)
  (for/sum ([b1 (in-list bl1)]
            [b2 (in-list bl2)])
       #:break (not (equal? b1 b2))
       1))

;; An RTresult is one of
;; (list 'exact RT) - An exact match for the given key that landed on the returned RT
;; (list 'partial RT RTedge Nat RTkey [Listof RT]) - A partial match stopping at the given
;;     RT, next partially matched edge, length of partial match, tail of the key not exactly
;;     matched, and a list of visited nodes ordered by most recent first
;; (list 'no-match RT RTkey [Listof RT]) - No match was found, but we stopped at the given
;;     RT, with tail of the key that did not match, and list of visited nodes ordered by most
;;     recent first

;; RT RTkey [[Listof RT]] -> RTresult
;; Helper function that performs tree iteration
(define (rt-partial-iterate node key [visited '()])
  (cond
    [(empty? key) (list 'exact node (cons node visited))]
    [else
     (let* ([bit (first key)]
            [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 (bl-common-len next-label key))])
       (cond
         [(and next-edge (= next-common-len (length next-label)))
          (rt-partial-iterate next-target (drop key next-common-len) (cons node visited))]
         [next-edge (list 'partial node next-edge next-common-len key (cons node visited))]
         [else (list 'no-match node key (cons node visited))]))]))

;; RT RTkey Any -> Void
;; Updates the radix tree for the given key, adding the elem to the data at that key.
(define (rt-add! node key elem)
  (define (insert-node! node key)
    (let* ([bit (first key)]
           [setter! (rt-setter bit)])
      (setter! node (rt-edge key (rt-node #f #f (list elem))))))
  (define (split-node! node key orig-edge prefix-len)
    (let* ([bit (first key)]
           [setter! (rt-setter bit)]
           [orig-label (rt-edge-label orig-edge)]
           [orig-target (rt-edge-target orig-edge)]
           [new-orig-edge (rt-edge (drop orig-label prefix-len) orig-target)])
      (if (= prefix-len (length key))
        (let* ([next-bit (list-ref orig-label prefix-len)]
               [common-node (rt-node (if next-bit #f new-orig-edge)
                                     (if next-bit new-orig-edge #f)
                                     (list elem))]
               [common-edge (rt-edge (take key prefix-len) common-node)])
          (setter! node common-edge))
        (let* ([new-insert-edge (rt-edge (drop key prefix-len)
                                         (rt-node #f #f (list elem)))]
               [diff-bit (list-ref key prefix-len)]
               [common-node (rt-node (if diff-bit new-orig-edge new-insert-edge)
                                     (if diff-bit new-insert-edge new-orig-edge)
                                     '())]
               [common-edge (rt-edge (take key prefix-len) common-node)])
          (setter! node common-edge)))))
  (match (rt-partial-iterate node key)
    [(list 'exact node _)
     (set-rt-node-data! node
                        (cons elem (rt-node-data node)))]
    [(list 'partial node orig-edge prefix-len partial-key visited)
     (split-node! node partial-key orig-edge prefix-len)]
    [(list 'no-match node partial-key visited)
     (insert-node! node partial-key)]))

;; RT RTKey (Any -> Boolean) -> Void
;; Delete all entries under the given key
(define (rt-del! node key del?)
  (for ([node (in-list (match (rt-partial-iterate node key)
                         [(list 'exact node _) (list node)]
                         [(list _ ... visited) visited]))])
    (set-rt-node-data! node
                       (filter (negate del?)
                               (rt-node-data node)))))

;; RT RTkey -> [Listof Any]
;; Looks up all the matches for the given key in the tree, ordered from most-specific to
;; least-specific.
(define (rt-lookup node key)
  (match (rt-partial-iterate node key)
    [(list _ ... visited) (append-map rt-node-data visited)]))

;; RT -> [Listof Any]
;; Converts the tree into a flat list of all contained node data
(define (rt-flatten node)
  (define (flatten-edge e)
    (if e (rt-flatten (rt-edge-target e)) '()))
  (append (rt-node-data node)
          (flatten-edge (rt-node-edge0 node))
          (flatten-edge (rt-node-edge1 node))))

;; RT [Str] ->
;; Debug print function that dumps the tree to current-output-port in a vaguely human-readable
;; format
(define (rt-dump node [prefix ""])
  (displayln (format "~a node ~s" 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 (rt-edge-label edge0)))
    (rt-dump (rt-edge-target edge0) (string-append prefix "  ")))
  (when edge1
    (displayln (format "~a  edge1 ~a" prefix (rt-edge-label edge1)))
    (rt-dump (rt-edge-target edge1) (string-append prefix "  "))))


(module+ test

  (define (str->bl x)
    (map (curry equal? #\1) (string->list x)))

  (define test (make-rt))

  (define (test-insert! x)
    (rt-add! test (str->bl x) x))

  (test-insert! "0001")
  (test-insert! "")
  (rt-dump test)
  (test-insert! "0001")
  (test-insert! "000")
  (test-insert! "0")
  (rt-dump test)
  (rt-partial-iterate test (str->bl "0001"))
  (rt-lookup test (str->bl "0000"))

  ; (test-insert! "0001")
  ; (test-insert! "1000")
  ; (test-insert! "1010")
  ; (test-insert! "0011")
  ; (test-insert! "0000")
  )