bvh2.lisp

(in-package #:org.shirakumo.fraf.trial.space.bvh2)

;;; CF https://www.researchgate.net/publication/254007711_Fast_Effective_BVH_Updates_for_Animated_Scenes

;;; This code uses the VEC4 type to represent 2D rectangles,
;;; indirectly by including VEC4 in the QUADTREE-NODE structure and
;;; also directly. In this representation, the vector components have
;;; the following semantics:
;;; x: x coordinate of top left corner
;;; x: y coordinate of top left corner
;;; z: x coordinate of bottom right corner
;;; w: y coordinate of bottom right corner

(defstruct (bvh-node
            (:include vec4)
            (:constructor %make-bvh-node (org.shirakumo.fraf.math.vectors::varr4 d p l r o))
            (:copier NIL)
            (:predicate NIL))
  (d 0 :type (unsigned-byte 16))
  (p NIL :type (or null bvh-node))
  (l NIL :type (or null bvh-node))
  (r NIL :type (or null bvh-node))
  (o NIL :type T))

(defmethod print-object ((node bvh-node) stream)
  (print-unreadable-object (node stream :type T)
    (let ((o (when (bvh-node-o node) (princ-to-string (bvh-node-o node)))))
      (format stream "~f ~f ~f ~f ~d~@[ ~a~]" (vx4 node) (vy4 node) (vz4 node) (vw4 node) (bvh-node-d node) o))))

(defun make-bvh-node-for (object parent)
  (let ((loc (location object))
        (siz (bsize object))
        (arr (make-array 4 :element-type 'single-float)))
    (setf (aref arr 0) (- (vx loc) (vx siz)))
    (setf (aref arr 1) (- (vy loc) (vy siz)))
    (setf (aref arr 2) (+ (vx loc) (vx siz)))
    (setf (aref arr 3) (+ (vy loc) (vy siz)))
    (%make-bvh-node arr (1+ (bvh-node-d parent)) parent NIL NIL object)))

(defun node-refit-object (node object)
  (let ((x (vx4 node)) (y (vy4 node))
        (z (vz4 node)) (w (vw4 node))
        (loc (location object))
        (siz (bsize object)))
    (setf (vx4 node) (- (vx loc) (vx siz)))
    (setf (vy4 node) (- (vy loc) (vy siz)))
    (setf (vz4 node) (+ (vx loc) (vx siz)))
    (setf (vw4 node) (+ (vy loc) (vy siz)))
    (when (or (/= x (vx4 node))
              (/= y (vy4 node))
              (/= z (vz4 node))
              (/= w (vw4 node)))
      (when (bvh-node-p node)
        (node-refit (bvh-node-p node)))
      T)))

(defun node-fit (node l r)
  (setf (vx4 node) (min (vx4 l) (vx4 r)))
  (setf (vy4 node) (min (vy4 l) (vy4 r)))
  (setf (vz4 node) (max (vz4 l) (vz4 r)))
  (setf (vw4 node) (max (vw4 l) (vw4 r))))

(defun node-refit (node)
  (let ((x (vx4 node)) (y (vy4 node))
        (z (vz4 node)) (w (vw4 node)))
    (node-fit node (bvh-node-l node) (bvh-node-r node))
    (when (or (/= x (vx4 node))
              (/= y (vy4 node))
              (/= z (vz4 node))
              (/= w (vw4 node)))
      (when (bvh-node-p node)
        (node-refit (bvh-node-p node)))
      T)))

(defun node-split (node object)
  (let ((l (make-bvh-node-for (bvh-node-o node) node))
        (r (make-bvh-node-for object node)))
    (setf (bvh-node-o node) NIL)
    (setf (bvh-node-l node) l)
    (setf (bvh-node-r node) r)
    (node-refit node)
    node))

(declaim (inline node-contains-p*))
(defun node-contains-p* (node region)
  (declare (optimize speed (safety 0)))
  (declare (type vec4 node region))
  (let ((nx (vx4 node))
        (ny (vy4 node))
        (nz (vz4 node))
        (nw (vw4 node))
        (rx (vx4 region))
        (ry (vy4 region))
        (rz (vz4 region))
        (rw (vw4 region)))
    (and (<= nx rx)
         (<= ny ry)
         (<= rz nz)
         (<= rw nw))))

(defun node-contains-p (node object)
  (declare (optimize speed (safety 0)))
  (declare (type vec4 node))
  (let* ((center (location object))
         (size/2 (bsize object))
         (region (vec (- (vx center) (vx size/2))
                      (- (vy center) (vy size/2))
                      (+ (vx center) (vx size/2))
                      (+ (vy center) (vy size/2)))))
    (declare (type vec2 center size/2))
    (node-contains-p* node region)))

(declaim (inline node-overlaps-p))
(defun node-overlaps-p (node region)
  (declare (optimize speed (safety 0)))
  (declare (type vec4 node region))
  (let ((nx (vx4 node))
        (ny (vy4 node))
        (nz (vz4 node))
        (nw (vw4 node))
        (rx (vx4 region))
        (ry (vy4 region))
        (rz (vz4 region))
        (rw (vw4 region)))
    (and (<= nx rz)
         (<= ny rw)
         (<= rx nz)
         (<= ry nw))))

(defun node-sub-p (node sub)
  (declare (optimize speed))
  (and (<= (vx4 node) (vx4 sub))
       (<= (vy4 node) (vy4 sub))
       (<= (vz4 sub) (vz4 node))
       (<= (vw4 sub) (vw4 node))))

(defun better-fit (a b object)
  (declare (optimize speed))
  (let ((ca (node-contains-p a object))
        (cb (node-contains-p b object)))
    (cond ((eq ca cb)
           ;; If it's in neither or in both, just see whose centroid we're closer to.
           (let ((ax (/ (+ (vz4 a) (vx4 a)) 2.0))
                 (ay (/ (+ (vw4 a) (vy4 a)) 2.0))
                 (bx (/ (+ (vz4 b) (vx4 b)) 2.0))
                 (by (/ (+ (vw4 b) (vy4 b)) 2.0))
                 (loc (location object)))
             (if (< (+ (expt (- ax (vx loc)) 2)
                       (expt (- ay (vy loc)) 2))
                    (+ (expt (- bx (vx loc)) 2)
                       (expt (- by (vy loc)) 2)))
                 a b)))
          (ca a)
          (cb b))))

(defun node-insert (node object)
  (declare (optimize speed))
  (cond ((bvh-node-o node)
         (node-split node object))
        ((bvh-node-r node)
         (node-insert (better-fit (bvh-node-l node) (bvh-node-r node) object) object))
        (T
         (setf (bvh-node-o node) object)
         (node-refit-object node object)
         node)))

(defun node-sibling (node)
  (declare (optimize speed))
  (let* ((p (bvh-node-p node))
         (l (bvh-node-l p))
         (r (bvh-node-r p)))
    (cond ((eq node l) r)
          ((eq node r) l)
          (T (error "What the fuck?")))))

(defun set-depth (node d)
  (declare (optimize speed))
  (setf (bvh-node-d node) d)
  (unless (bvh-node-o node)
    (set-depth (bvh-node-l node) (1+ d))
    (set-depth (bvh-node-r node) (1+ d))))

(defun node-transfer (target source)
  (declare (optimize speed))
  (setf (vx4 target) (vx4 source))
  (setf (vy4 target) (vy4 source))
  (setf (vz4 target) (vz4 source))
  (setf (vw4 target) (vw4 source))
  (let ((l (bvh-node-l source))
        (r (bvh-node-r source)))
    (setf (bvh-node-l target) l)
    (setf (bvh-node-r target) r)
    (when l
      (setf (bvh-node-p l) target)
      (setf (bvh-node-p r) target))
    (setf (bvh-node-o target) (bvh-node-o source)))
  (set-depth target (bvh-node-d source)))

(defun node-remove (node)
  (declare (optimize speed))
  (let ((p (bvh-node-p node)))
    (cond (p
           (node-transfer p (node-sibling node))
           (when (bvh-node-p p)
             (node-refit (bvh-node-p p)))
           p)
          (T
           (setf (bvh-node-o node) NIL)
           node))))

(defstruct (bvh
            (:include container)
            (:constructor make-bvh ())
            (:copier NIL)
            (:predicate NIL))
  (root (%make-bvh-node (make-array 4 :element-type 'single-float :initial-element 0f0) 0 NIL NIL NIL NIL) :type bvh-node)
  (table (make-hash-table :test 'eq) :type hash-table))

(defun bvh-insert (bvh object)
  (declare (optimize speed))
  (let* ((table (bvh-table bvh))
         (existing (gethash object table)))
    (if existing
        (node-refit-object existing object)
        (let ((node (node-insert (bvh-root bvh) object)))
          (cond ((eq object (bvh-node-o node))
                 (setf (gethash object table) node))
                (T
                 (setf (gethash (bvh-node-o (bvh-node-l node)) table) (bvh-node-l node))
                 (setf (gethash (bvh-node-o (bvh-node-r node)) table) (bvh-node-r node))))))
    object))

(defun bvh-remove (bvh object)
  (declare (optimize speed))
  (let* ((table (bvh-table bvh))
         (node (gethash object table)))
    (when node
      (remhash object table)
      (let ((p (node-remove node)))
        (setf (gethash (bvh-node-o p) table) p)))))

(defun bvh-update (bvh object)
  (declare (optimize speed))
  ;; FIXME: Figure out when to rebalance the tree.
  (let ((node (gethash object (bvh-table bvh))))
    (when node
      (node-refit-object node object))))

(defmethod enter (object (bvh bvh))
  (bvh-insert bvh object))

(defmethod leave (object (bvh bvh))
  (bvh-remove bvh object))

(defmethod update (object (bvh bvh))
  (bvh-update bvh object))

(defmethod clear ((bvh bvh))
  (clrhash (bvh-table bvh))
  (setf (bvh-root bvh) (%make-bvh-node (make-array 4 :element-type 'single-float :initial-element 0f0) 0 NIL NIL NIL NIL))
  bvh)

(defmethod describe-object ((bvh bvh) stream)
  (call-next-method)
  (format stream "~&~%-------------------------")
  (org.shirakumo.text-draw:tree
   (bvh-root bvh)
   (lambda (node)
     (unless (bvh-node-o node)
       (list (bvh-node-l node) (bvh-node-r node))))
   :stream stream))

(defun bvh-lines (bvh)
  (let ((p ()))
    (labels ((depth-color (depth)
               (let ((d (max 0.0 (- 1.0 (/ depth 100)))))
                 (vec 1 d d 0.1)))
             (recurse (node)
               (let ((color (depth-color (bvh-node-d node))))
                 (push (list (vxy_ node) color) p)
                 (push (list (vzy_ node) color) p)
                 (push (list (vxw_ node) color) p)
                 (push (list (vzw_ node) color) p)
                 (push (list (vxy_ node) color) p)
                 (push (list (vxw_ node) color) p)
                 (push (list (vzy_ node) color) p)
                 (push (list (vzw_ node) color) p)
                 (when (bvh-node-l node)
                   (recurse (bvh-node-l node))
                   (recurse (bvh-node-r node))))))
      (recurse (bvh-root bvh)))
    p))

(defmethod check ((bvh bvh))
  (labels ((recurse (node)
             (cond ((bvh-node-l node)
                    (unless (eq node (bvh-node-p (bvh-node-l node)))
                      (error "The left child~%  ~a~%is not parented to~%  ~a"
                             (bvh-node-l node) node))
                    (unless (eq node (bvh-node-p (bvh-node-r node)))
                      (error "The right child~%  ~a~%is not parented to~%  ~a"
                             (bvh-node-r node) node))
                    (unless (node-sub-p node (bvh-node-l node))
                      (error "The parent node~%  ~a~%does not contain the left child~%  ~a"
                             node (bvh-node-l node)))
                    (unless (node-sub-p node (bvh-node-r node))
                      (error "The parent node~%  ~a~%does not contain the right child~%  ~a"
                             node (bvh-node-r node)))
                    (recurse (bvh-node-l node))
                    (recurse (bvh-node-r node)))
                   ((bvh-node-o node)
                    (unless (eq node (gethash (bvh-node-o node) (bvh-table bvh)))
                      (error "The node~%  ~a~%is not assigned to object~%  ~a~%as it is assigned to~%  ~a"
                             node (bvh-node-o node) (gethash (bvh-node-o node) (bvh-table bvh))))))))
    (recurse (bvh-root bvh)))
  (loop for o being the hash-keys of (bvh-table bvh)
        for n being the hash-values of (bvh-table bvh)
        do (unless (eq o (bvh-node-o n))
             (error "The node~%  ~a~%does not refer to object~%  ~a~%and instead tracks~%  ~a"
                    n o (bvh-node-o n)))))

(defun bvh-refit (bvh)
  (declare (optimize speed))
  (labels ((recurse (node)
             (cond ((bvh-node-l node)
                    (recurse (bvh-node-l node))
                    (recurse (bvh-node-r node))
                    (node-refit node))
                   (T
                    (node-refit-object node (bvh-node-o node))))))
    (recurse (bvh-root bvh))))

(defmethod call-with-all (function (bvh bvh))
  (labels ((recurse (node)
             (when node
               (let ((o (bvh-node-o node)))
                 (cond (o
                        (funcall function o))
                       (T
                        (recurse (bvh-node-l node))
                        (recurse (bvh-node-r node))))))))
    (recurse (bvh-root bvh))))

(defmethod call-with-contained (function (bvh bvh) (region region))
  (declare (optimize speed (safety 1)))
  (let* ((function (ensure-function function))
         (size (region-size region))
         (region (vec (vx3 region) (vy3 region)
                      (+ (vx3 region) (vx3 size)) (+ (vy3 region) (vy3 size))))
         (tentative (make-array 256))
         (i 0))
    (declare (type (integer 0 256) i))
    (declare (dynamic-extent tentative region))
    (flet ((add (node)
             (when (and (< i 256) node)
               (setf (svref tentative i) node)
               (incf i))))
      (declare (inline add))
      (add (bvh-root bvh))
      (loop (decf i)
            (let* ((node (svref tentative i))
                   (o (bvh-node-o node)))
              (cond (o
                     ;; TODO(jmoringe): should make a shared utility function
                     (when (and (<= (vx region) (vx node)) ; REGION contains NODE
                                (<= (vy region) (vy node))
                                (<= (vz node) (vz region))
                                (<= (vw node) (vw region)))
                       (funcall function o)))
                    (T
                     (when (node-overlaps-p node region)
                       (add (bvh-node-l node))
                       (add (bvh-node-r node)))))
              (when (= 0 i)
                (return)))))))

(defmethod call-with-overlapping (function (bvh bvh) (region region))
  (declare (optimize speed (safety 1)))
  (let* ((function (ensure-function function))
         (size (region-size region))
         (region (vec (vx3 region) (vy3 region)
                      (+ (vx3 region) (vx3 size)) (+ (vy3 region) (vy3 size)))))
    (declare (dynamic-extent region))
    (labels ((recurse (node)
               (when (and node (node-overlaps-p node region))
                 (let ((o (bvh-node-o node)))
                   (cond (o
                          (funcall function o))
                         (T
                          (recurse (bvh-node-l node))
                          (recurse (bvh-node-r node))))))))
      (recurse (bvh-root bvh)))))

(defstruct (bvh-iterator
            (:constructor make-bvh-iterator (bvh region))
            (:copier NIL)
            (:predicate NIL))
  (bvh NIL :type bvh)
  (region NIL :type vec4))

(defmethod for:make-iterator ((bvh bvh) &key in)
  (if in
      (make-bvh-iterator bvh in)
      bvh))

(defmethod for:step-functions ((iterator bvh-iterator))
  (declare (optimize speed))
  (let ((node (bvh-root (bvh-iterator-bvh iterator)))
        (region (bvh-iterator-region iterator)))
    (labels ((next-leaf (node child)
               (when node
                 (let ((l (bvh-node-l node))
                       (r (bvh-node-r node)))
                   (cond ((bvh-node-o node)
                          node)
                         ((null child)
                          (if (node-overlaps-p l region)
                              (next-leaf l NIL)
                              (next-leaf node l)))
                         ((eq child l)
                          (if (node-overlaps-p r region)
                              (next-leaf r NIL)
                              (next-leaf node r)))
                         ((eq child r)
                          (next-leaf (bvh-node-p node) node)))))))
      (setf node (next-leaf node NIL))
      (values
       (lambda ()
         (prog1 (bvh-node-o node)
           (setf node (next-leaf (bvh-node-p node) node))))
       (lambda ()
         node)
       (lambda (value)
         (declare (ignore value))
         (error "Not supported"))
       (lambda ())))))

(defmethod for:step-functions ((bvh bvh))
  (declare (optimize speed))
  (let ((node (bvh-root bvh)))
    (labels ((next-leaf (node child)
               (when node
                 (cond ((bvh-node-o node)
                        node)
                       ((null child)
                        (next-leaf (bvh-node-l node) NIL))
                       ((eq child (bvh-node-l node))
                        (next-leaf (bvh-node-r node) NIL))
                       ((eq child (bvh-node-r node))
                        (next-leaf (bvh-node-p node) node))))))
      (setf node (next-leaf node NIL))
      (values
       (lambda ()
         (prog1 (bvh-node-o node)
           (setf node (next-leaf (bvh-node-p node) node))))
       (lambda ()
         node)
       (lambda (value)
         (declare (ignore value))
         (error "Not supported"))
       (lambda ())))))

(declaim (inline perimiter-heuristic cost node-cost node-priority))

(defun perimiter-heuristic (x y z w)
  (declare (optimize speed (safety 0))
           (type single-float x y z w))
  (+ (- z x)
     (- w y)))

(defun cost (x y z w)
  (declare (optimize speed (safety 0)))
  (the single-float
       (perimiter-heuristic x y z w)))

(defun node-cost (node)
  (declare (optimize speed (safety 0))
           (type vec4 node))
  (the single-float
       (cost (vx4 node) (vy4 node)
             (vz4 node) (vw4 node))))

(defun node-priority (node)
  (declare (optimize speed (safety 0)))
  (if (bvh-node-p node)
      (/ (node-cost (bvh-node-p node))
         (node-cost node))
      0.0))

;; like "Parallel Reinsertion for Bounding Volume Hierarchy Optimization" (D. Meister and J. Bittner)
;; the paper is mostly about optimization for paralellism, with the algorithm being derived from
;; "Fast Insertion-Based Optimization of Bounding Volume Hierarchies" (J. Bittner, M. Hapala, and V. Havran)
;;
;; Returns the node for which NODE being reinserted as its child would
;; create the greatest cumulative decrease in box perimiter.
(defun node-find-best-reinsertion-position (node bvh)
  (declare (optimize speed))
  (when (bvh-node-p node)
    (let* ((parent (bvh-node-p node))
           (pivot parent)
           (pivot-box (vec4 0.0 0.0 0.0 0.0))
           (pivot-box-initialized NIL)
           (decrease-bound (- (node-cost parent) (node-cost node)))
           (decrease 0.0f0)
           (decrease-best decrease)
           (out (node-sibling node))
           (out-best out)
           (down T))
      (declare (dynamic-extent pivot-box)
               (type single-float decrease-bound decrease decrease-best))
      (flet ((union-cost (node-a node-b)
               (declare (optimize speed (safety 0)))
               (cost (min (vx4 node-a)
                          (vx4 node-b))
                     (min (vy4 node-a)
                          (vy4 node-b))
                     (max (vz4 node-a)
                          (vz4 node-b))
                     (max (vw4 node-a)
                          (vw4 node-b)))))
        (declare (inline union-cost))
        (loop (cond (down
                     ;; traverse downwards
                     (let* ((merged-cost (union-cost node out))
                            ;; how much better off we would be if OUT was NODE's parent
                            (decrease-direct (- (node-cost parent)
                                                merged-cost)))
                       (when (< decrease-best (+ decrease decrease-direct))
                         (setf decrease-best (+ decrease decrease-direct))
                         (setf out-best out))
                       ;; decrease of cost gets worse as you descend
                       (decf decrease (- merged-cost (node-cost out)))
                       (if (or (bvh-node-o out) ; stop at leaves
                               (<= (+ decrease-bound decrease)
                                   ;; early stop heuristic; if descending
                                   ;; would be worse than where we started,
                                   ;; stop descending.
                                   decrease-best))
                           (setf down NIL)
                           (setf out (bvh-node-l out)))))
                    (T ; traverse upwards towards pivot
                     ;; decrease of cost improves as you ascend
                     (incf decrease (- (union-cost node out) (node-cost out)))
                     (cond ((eq pivot (bvh-node-p out))        ; we've reached the pivot
                            (cond (pivot-box-initialized ; expand the pivot box to include OUT
                                   (node-fit pivot-box pivot-box out))
                                  (T
                                   (vsetf pivot-box (vx out) (vy out) (vz out) (vw out))
                                   (setf pivot-box-initialized T)))
                            (setf out (bvh-node-p out)) ; consider the pivot as OUT,
                            (when (not (eq out parent)) ; but not if we're still on the original parent.
                              ;; how much better off we would be if PIVOT-BOX was NODE's parent
                              (let ((decrease-direct (- (node-cost parent)
                                                        (union-cost node pivot-box))))
                                (when (< decrease-best (+ decrease decrease-direct))
                                  (setf decrease-best (+ decrease decrease-direct))
                                  ;; set pivot as best
                                  (setf out-best out)))
                              (decf decrease (- (node-cost pivot-box)
                                                (node-cost out))))
                            (when (eq out (bvh-root bvh)) ; stop at root
                              ;; If the best position found is the
                              ;; sibling, NODE is already in the best spot.
                              (return (unless (eq out-best (node-sibling node))
                                        out-best)))
                            ;; when we haven't stopped, move the pivot up
                            ;; and traverse down on the sibling branch
                            (setf out (node-sibling pivot))
                            (setf pivot (bvh-node-p out))
                            (setf down T))

                         ;; we haven't reached the pivot. traverse sibling down, or parent up
                         ((eq out (bvh-node-l (bvh-node-p out)))
                          (setf down T)
                          (setf out (node-sibling out)))
                         (T
                          (setf out (bvh-node-p out)))))))))))

;; Remove NODE from BVH, then splice it back into the tree so that OUT is its parent.
(defun node-remove-and-reinsert (bvh node out)
  (declare (optimize speed))
  (let ((sibling (node-sibling node))
        (parent (bvh-node-p node)))
    (when parent
      ;; Put SIBLING where PARENT is
      (let ((grandparent (bvh-node-p parent)))
        (cond (grandparent
               (if (eq (bvh-node-l grandparent) parent)
                   (setf (bvh-node-l grandparent) sibling)
                   (setf (bvh-node-r grandparent) sibling))
               (setf (bvh-node-p sibling) grandparent))
              (T ;; if there's no grandparent, parent is the root.
               (setf (bvh-root bvh) sibling)
               (setf (bvh-node-p sibling) NIL)))
        (set-depth sibling (1- (bvh-node-d sibling)))
        (when grandparent
          (node-refit grandparent)))
      ;; Put PARENT where OUT is, with NODE and OUT as its children
      (let ((grandparent (bvh-node-p out)))
        (cond (grandparent
               (if (eq (bvh-node-l grandparent) out)
                   (setf (bvh-node-l grandparent) parent)
                   (setf (bvh-node-r grandparent) parent))
               (setf (bvh-node-p parent) grandparent))
              (T ;; out is the root
               (setf (bvh-root bvh) parent)
               (setf (bvh-node-p parent) NIL)))
        (setf (bvh-node-r parent) out)
        (setf (bvh-node-l parent) node)
        (setf (bvh-node-p out) parent)
        (setf (bvh-node-p node) parent)
        (set-depth parent (bvh-node-d out))
        (node-refit parent)))))

(defun node-reinsert-to-best (node bvh)
  (let ((best (node-find-best-reinsertion-position node bvh)))
    (when best
      (node-remove-and-reinsert bvh node best)
      T)))

(defmethod reoptimize ((bvh bvh) &key (rounds 1))
  (dotimes (i rounds bvh)
    (labels ((descend (node)
               (node-reinsert-to-best node bvh)
               (unless (bvh-node-o node)
                 (descend (bvh-node-l node))
                 (descend (bvh-node-r node))))
             (bvh-empty-p (bvh)
               (let ((root (bvh-root bvh)))
                 (not (or (bvh-node-o root)
                          (bvh-node-l root)
                          (bvh-node-r root))))))
      (unless (bvh-empty-p bvh)
        (descend (bvh-root bvh))))))