;;;; -*- Mode: Lisp; Base: 10; Syntax: ANSI-Common-Lisp; Package: kira; Coding: utf-8 -*-
;;;; Copyright © 2024 David Mullen. All Rights Reserved. Origin: <https://cl-pdx.com/kira/>

(in-package :kira)

(defgeneric ternary-node-priority (generic-node)
  (:method ((ternary-node-as-list list)) 0))

(define-persistent-class ternary-node ()
  ((label :initarg :label :accessor ternary-node-label)
   (data :initform nil :initarg :data :accessor ternary-node-data)
   (left :initform nil :initarg :left :accessor ternary-node-left)
   (middle :initform nil :initarg :middle :accessor ternary-node-middle)
   (right :initform nil :initarg :right :accessor ternary-node-right)
   (priority :initform 0 :initarg :priority :accessor ternary-node-priority)))

(defconstant +max-priority+ (1- (ash 1 55)))
(defconstant +ternary-burst-threshold+ 64)

(defun ternary-search (key position node)
  (let ((current-label (char key position)))
    (cond ((listp node)
           (loop for (entry-key . entry-key-association) in node
                 if (string-equal entry-key key :start2 position)
                 return entry-key-association))
          ((char-equal (ternary-node-label node) current-label)
           (if (= (the array-index (1+ position)) (length key))
               (ternary-node-data node)
               (let ((middle (ternary-node-middle node)))
                 (ternary-search key (1+ position) middle))))
          ((char-greaterp current-label (ternary-node-label node))
           (ternary-search key position (ternary-node-right node)))
          ((ternary-search key position (ternary-node-left node))))))

(defun get-term (key ternary-search-tree)
  (ternary-search (string key) 0 ternary-search-tree))

(define-compiler-macro get-term (key ternary-search-tree)
  `(ternary-search (string ,key) 0 ,ternary-search-tree))

(defun ternary-rotate-with-left (node)
  (let ((left (ternary-node-left node)))
    (with-accessors ((place ternary-node-right)) left
      (setf (ternary-node-left node) place place node)
      left)))

(defun ternary-rotate-with-right (node)
  (let ((right (ternary-node-right node)))
    (with-accessors ((place ternary-node-left)) right
      (setf (ternary-node-right node) place place node)
      right)))

(define-setf-expander get-term (key place &environment environment)
  (multiple-value-bind (temps vals news writer reader) (get-setf-expansion place environment)
    (unless (endp (cdr news)) (error "Multiple values aren't supported in ternary search trees."))
    (with-gensyms (key-string node new-data)
      (values `(,key-string ,@temps ,node)
              `((string ,key) ,@vals ,reader)
              `(,new-data)
              `(prog1 ,new-data
                 (%ternary-insert ,key-string ,new-data 0 ,node)
                 (when (or (listp ,node) (not (eq ,node ,reader)))
                   (let ((,(car news) ,node)) ,writer)))
              `(get-term ,key-string ,reader)))))

(defun nlabeled (split-key alist &aux (new-ternary-node-label (char split-key 0)))
  "Destructively strip the leading character from each key in ALIST matching SPLIT-KEY."
  (loop for entry in alist if (char-equal (char (car entry) 0) new-ternary-node-label)
        collect entry into middle and do (setf (car entry) (subseq (car entry) 1))
        finally (return (values new-ternary-node-label middle))))

(defun alist-insert (key data position alist)
  (assert (< position (length key)) (key position) "Key ~S exhausted." key)
  (loop for entry in alist if (string-equal (car entry) key :start2 position)
        do (return-from alist-insert (progn (setf (cdr entry) data) alist))
        finally (let ((alist-delta (list (cons (subseq key position) data))))
                  (setq alist (merge 'list alist alist-delta #'string-lessp :key #'car))
                  (let ((list-under-threshold-p (< (length alist) +ternary-burst-threshold+)))
                    (when list-under-threshold-p (return-from alist-insert alist))
                    (let* ((midpoint (floor (length alist) 2))
                           (split-key (car (nth midpoint alist))))
                      (multiple-value-bind (label middle) (nlabeled split-key alist)
                        (let* ((priority (1+ (the binary-fixnum (random +max-priority+))))
                               (alist (sort (set-difference alist middle) #'string-lessp :key #'car))
                               (right-position (position label alist :test #'string-lessp :key #'car))
                               (data (whereas ((empty (assoc "" middle :test #'string=))) (cdr empty)))
                               (right (when right-position (nthcdr right-position alist)))
                               (left (if right (ldiff alist right) alist)))
                          (deletef "" middle :key #'car :test #'string=)
                          (return (make-instance 'ternary-node
                                                 :label label :left left
                                                 :data data :priority priority
                                                 :middle middle :right right)))))))))

(defmacro set-if-changed (place new-value &environment environment)
  (multiple-value-bind (temps vals news writer reader) (get-setf-expansion place environment)
    (unless (endp (cdr news)) (error "Multiple values aren't supported in SET-IF-CHANGED."))
    (with-gensyms (old-temp new-temp)
      `(let* (,.(mapcar #'list temps vals) (,old-temp ,reader) (,new-temp ,new-value))
         (if (eq ,new-temp ,old-temp) ,old-temp (let ((,(car news) ,new-temp)) ,writer))))))

(defmacro %ternary-insert (key data position place)
  `(cond ((listp ,place) (setq ,place (alist-insert ,key ,data ,position ,place)))
         (t (set-if-changed ,place (ternary-insert ,key ,data ,position ,place)))))

(defun ternary-key-label (key position)
  (unless (< position (length key))
    (error "Key ~S exhausted." key))
  (char key position))

(defun ternary-insert (key data position node)
  (let ((current-label (ternary-key-label key position)))
    (with-slots* (priority left middle right label) node
      (cond ((char-equal current-label label)
             (if (= (the array-index (1+ position)) (length key))
                 (progn (setf (ternary-node-data node) data) node)
                 (progn (%ternary-insert key data (1+ position) middle)
                        (with-accessors ((middle-priority ternary-node-priority)) middle
                          (when (> middle-priority priority) (setq priority middle-priority))
                          node))))
            ((char-greaterp current-label label)
             (%ternary-insert key data position right)
             (let ((rotatep (> (ternary-node-priority right) priority)))
               (if rotatep (ternary-rotate-with-right node) node)))
            ((%ternary-insert key data position left)
             (let ((rotatep (> (ternary-node-priority left) priority)))
               (if rotatep (ternary-rotate-with-left node) node)))))))

(defmacro delete-term (key place &environment environment)
  (multiple-value-bind (temps vals news writer reader) (get-setf-expansion place environment)
    (unless (endp (cdr news)) (error "Multiple values aren't supported in ternary search trees."))
    (let ((temporary-bindings (mapcar #'list temps vals)) (new-place-store-variable (car news)))
      (with-gensyms (key-temp place-temp)
        `(let* ((,key-temp (string ,key))
                ,.temporary-bindings
                (,place-temp ,reader))
           (%ternary-delete ,key-temp 0 ,place-temp)
           (when (or (listp ,place-temp) (not (eq ,place-temp ,reader)))
             (let ((,new-place-store-variable ,place-temp)) ,writer))
           ,place-temp)))))

(defmacro %ternary-delete (key position place &aux (k `(subseq ,key ,position)))
  `(cond ((listp ,place) (deletef ,k ,place :test #'string-equal :key #'car))
         (t (set-if-changed ,place (ternary-delete ,key ,position ,place)))))

(defun ternary-delete (key position node)
  (let ((current-label (ternary-key-label key position)))
    (with-slots* (priority left middle right label) node
      (if (char-equal current-label label)
          (cond ((< (the array-index (1+ position)) (length key))
                 (let ((new-middle (%ternary-delete key (the array-index (1+ position)) middle)))
                   (when (eq (or (ternary-node-data node) new-middle) nil) (setq priority 0)))
                 (ternary-heapify-or-delete node))
                ((progn (setf (ternary-node-data node) nil) middle) node)
                ((progn (setq priority 0) (ternary-heapify-or-delete node))))
          (prog1 node
            (if (char-greaterp current-label label)
                (%ternary-delete key position right)
                (%ternary-delete key position left)))))))

(defun ternary-heapify-or-delete (node)
  (with-slots* (priority left middle right data) node
    (unless (or left middle right data) (return-from ternary-heapify-or-delete))
    (with-accessors ((left-priority ternary-node-priority)) left
      (with-accessors ((right-priority ternary-node-priority)) right
        (cond ((>= priority (max right-priority left-priority)) node)
              ((> right-priority left-priority)
               (let ((rotated-in-node (ternary-rotate-with-right node)))
                 (with-accessors ((left ternary-node-left)) rotated-in-node
                   (set-if-changed left (ternary-heapify-or-delete left)))
                 rotated-in-node))
              ((let ((rotated-in-node (ternary-rotate-with-left node)))
                 (with-accessors ((right ternary-node-right)) rotated-in-node
                   (set-if-changed right (ternary-heapify-or-delete right)))
                 rotated-in-node)))))))

(defun map-terms (function ternary-search-tree)
  (cond ((typep ternary-search-tree 'ternary-node)
         (map-terms function (ternary-node-left ternary-search-tree))
         (whereas ((data (ternary-node-data ternary-search-tree)))
           (funcall function data))
         (map-terms function (ternary-node-middle ternary-search-tree))
         (map-terms function (ternary-node-right ternary-search-tree))
         ternary-search-tree)
        ((loop for entry in ternary-search-tree
               do (funcall function (cdr entry))
               finally (return ternary-search-tree)))))

(defmacro do-terms ((var ternary-search-tree &optional result-form) &body body)
  `(block nil (map-terms #'(lambda (,var) (tagbody ,@body)) ,ternary-search-tree)
              ,result-form))

(defgeneric xvector-depth (xvector)
  (:method ((null-xv null)) 0))

(defgeneric xvector-root (xvector)
  (:method ((null-xv null)) nil))

(defgeneric xvector-max (xvector)
  (:method ((null-xv null)) -1))

(define-persistent-class xvector ()
  ((depth :initform 0 :accessor xvector-depth)
   (root :initform nil :accessor xvector-root)
   (max :initform -1 :accessor xvector-max)))

(define-symbol-macro +base32hex+
  "0123456789ABCDEFGHIJKLMNOPQRSTUV")

(defun key-depth (key)
  (check-type key (integer 0 *) "a vector index")
  (1- (max (ceiling (integer-length key) 5) 1)))

(defun make-string-key (key &optional (depth (key-depth key)))
  (loop with string-key of-type simple-string = (make-string (1+ depth))
        for position of-type fixnum downfrom (* 5 depth) downto 0 by 5
        for val of-type (unsigned-byte 5) = (ldb (byte 5 position) key)
        for i of-type array-index from 0 to (the array-index depth)
        do (setf (schar string-key i) (schar +base32hex+ val))
        finally (return string-key)))

(defun xvref (xvector key)
  (unless xvector (return-from xvref))
  (let* ((root (xvector-root xvector))
         (depth (xvector-depth xvector))
         (k (make-string-key key depth)))
    (get-term k root)))

(define-setf-expander xvref (place key &environment environment)
  (multiple-value-bind (temps vals news writer reader) (get-setf-expansion place environment)
    (unless (endp (cdr news)) (error "Multiple values aren't supported by XVECTOR places."))
    (with-gensyms (key-temp xvector new-data)
      (values `(,key-temp ,@temps ,xvector)
              `(,key ,@vals ,reader)
              `(,new-data)
              `(progn
                 (unless ,xvector
                   (setq ,xvector (make-instance 'xvector))
                   (let ((,(car news) ,xvector)) ,writer))
                 (%xvector-set ,xvector ,key-temp ,new-data))
              `(xvref ,xvector ,key-temp)))))

(defun %xvector-set (xvector key data)
  (loop with root = (xvector-root xvector)
        with key-depth of-type fixnum = (key-depth key)
        and root-priority = (ternary-node-priority root)
        while (> key-depth (xvector-depth xvector))
        do (setf (xvector-root xvector)
                 (make-instance
                  'ternary-node :label #\0 :middle (xvector-root xvector)
                  :priority (maxf root-priority (1+ (random +max-priority+))))
                 (xvector-depth xvector)
                 ;; Extend by one key digit.
                 (1+ (xvector-depth xvector)))
        ;; Now the XVECTOR-DEPTH >= KEY-DEPTH.
        ;; Also adjust XVECTOR-MAX accordingly.
        finally (let ((string-key (make-string-key key (xvector-depth xvector))))
                  (setf (get-term string-key (xvector-root xvector)) data)
                  (unless (eq data nil) (maxf (xvector-max xvector) key))
                  (return data))))

(eval-always
  (defconstant +bitmap-burst-threshold+ 64)
  (defconstant +bitmap-bit-vector-size+ 14)
  (defconstant +bitmap-quad-space-size+ 16))

(defmacro %set-bit (part bit-index)
  `(let ((sub-index (ldb (byte +bitmap-bit-vector-size+ 0) ,bit-index)))
     (declare (type (unsigned-byte ,+bitmap-bit-vector-size+) sub-index))
     (prog1 ,part (setf (bit ,part sub-index) 1))))

(defun qxr (quad slot-index)
  "Get Q{A,B,C,D}R by index."
  (ecase slot-index
    (0 (qar quad))
    (1 (qbr quad))
    (2 (qcr quad))
    (3 (qdr quad))))

(defun (setf qxr) (new-value quad slot-index)
  "Set the slot value Q{A,B,C,D}R by index."
  (ecase slot-index
    (0 (setf (qar quad) new-value))
    (1 (setf (qbr quad) new-value))
    (2 (setf (qcr quad) new-value))
    (3 (setf (qdr quad) new-value))))

(defun %q-set-bit (quad bit-index)
  (let ((part-index (ldb (byte 2 +bitmap-bit-vector-size+) bit-index)))
    (flet ((make-bitmap-part () (make-bit-vector (expt 2 +bitmap-bit-vector-size+))))
      (let ((part (or (qxr quad part-index) (make-bitmap-part))))
        (%set-bit (setf (qxr quad part-index) part) bit-index)))))

(defun %xv-set-bit (bitmap i)
  (let* ((item-index (ash i (- +bitmap-quad-space-size+)))
         (existing-item (xvref bitmap item-index)))
    (multiple-value-prog1 (values bitmap item-index)
      (cond ((eql existing-item i)) ; No operation.
            ((not existing-item) (setf (xvref bitmap item-index) i))
            ((typep existing-item 'quad) (%q-set-bit existing-item i))
            ((symbol-macrolet ((item-at-index (xvref bitmap item-index)))
               (let ((quad (setq item-at-index (make-instance 'quad))))
                 (%q-set-bit quad existing-item) (%q-set-bit quad i))))))))

(defun %add-to-bitmap (bitmap i)
  (cond ((typep bitmap 'xvector) (%xv-set-bit bitmap i))
        (t (progn (pushnew i bitmap) (sortf bitmap #'<))
           (when (< (length bitmap) +bitmap-burst-threshold+)
             (return-from %add-to-bitmap bitmap))
           (let ((new (make-instance 'xvector)))
             (dolist (existing-bit-index bitmap new)
               (%xv-set-bit new existing-bit-index))))))

(defmacro add-to-bitmap (place new-bit-index &environment environment)
  (multiple-value-bind (temps vals news writer reader) (get-setf-expansion place environment)
    (unless (endp (cdr news)) (error "Multiple values aren't supported in ADD-TO-BITMAP."))
    (with-gensyms (bitmap new-bitmap)
      `(let* (,.(mapcar #'list temps vals) (,bitmap ,reader))
         (let ((,new-bitmap (%add-to-bitmap ,bitmap ,new-bit-index)))
           (when (or (listp ,bitmap) (not (eq ,new-bitmap ,bitmap)))
             (let ((,(car news) ,new-bitmap)) ,writer))
           ,new-bitmap)))))

(defun in-bitmap-p (bitmap i)
  (cond ((listp bitmap) (if (member i bitmap) t nil))
        ((whereas ((item (xvref bitmap (ash i (- +bitmap-quad-space-size+)))))
           (when (integerp item) (return-from in-bitmap-p (eql item i)))
           (let* ((part-index (ldb (byte 2 +bitmap-bit-vector-size+) i))
                  (quad-reader-function (svref #(qar qbr qcr qdr) part-index))
                  (index-in-part (ldb (byte +bitmap-bit-vector-size+ 0) i))
                  (part (funcall quad-reader-function item)))
             (declare (type (unsigned-byte 2) part-index))
             (declare (type array-index index-in-part))
             (and part (= (bit part index-in-part) 1)))))))