Add (chickadee data quadtree) module.

1 files changed, 338 insertions, 0 deletions

diff --git a/chickadee/data/quadtree.scm b/chickadee/data/quadtree.scm
new file mode 100644
index 0000000..731faf2
--- /dev/null
+++ b/chickadee/data/quadtree.scm
@@ -0,0 +1,338 @@
+;;; Chickadee Game Toolkit
+;;; Copyright © 2021 David Thompson <davet@gnu.org>
+;;;
+;;; Chickadee is free software: you can redistribute it and/or modify
+;;; it under the terms of the GNU General Public License as published
+;;; by the Free Software Foundation, either version 3 of the License,
+;;; or (at your option) any later version.
+;;;
+;;; Chickadee is distributed in the hope that it will be useful, but
+;;; WITHOUT ANY WARRANTY; without even the implied warranty of
+;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+;;; General Public License for more details.
+;;;
+;;; You should have received a copy of the GNU General Public License
+;;; along with this program.  If not, see
+;;; <http://www.gnu.org/licenses/>.
+
+(define-module (chickadee data quadtree)
+  #:use-module (chickadee math rect)
+  #:use-module (ice-9 format)
+  #:use-module (srfi srfi-9)
+  #:use-module (srfi srfi-9 gnu)
+  #:export (make-quadtree
+            quadtree?
+            quadtree-leaf?
+            quadtree-bounds
+            quadtree-max-depth
+            quadtree-max-size
+            quadtree-depth
+            quadtree-size
+            quadtree-q1
+            quadtree-q2
+            quadtree-q3
+            quadtree-q4
+            quadtree-clear!
+            quadtree-insert!
+            quadtree-delete!
+            quadtree-find
+            quadtree-fold
+            quadtree-for-each))
+
+;; The quadrants:
+;;
+;; *------*------*
+;; |      |      |
+;; |  Q2  |  Q1  |
+;; |      |      |
+;; *------*------*
+;; |      |      |
+;; |  Q3  |  Q4  |
+;; |      |      |
+;; *------*------*
+(define-record-type <quadtree>
+  (%make-quadtree bounds max-depth max-size depth size objects leaf?)
+  quadtree?
+  (bounds quadtree-bounds)
+  (max-depth quadtree-max-depth)
+  (max-size quadtree-max-size)
+  (depth quadtree-depth)
+  (size quadtree-size set-quadtree-size!)
+  (objects quadtree-objects set-quadtree-objects!)
+  (leaf? quadtree-leaf? set-quadtree-leaf!)
+  (q1 quadtree-q1 set-quadtree-q1!)
+  (q2 quadtree-q2 set-quadtree-q2!)
+  (q3 quadtree-q3 set-quadtree-q3!)
+  (q4 quadtree-q4 set-quadtree-q4!))
+
+(define (display-quadtree quadtree port)
+  (format port "#<quadtree bounds: ~s depth: ~d size: ~d>"
+          (quadtree-bounds quadtree)
+          (quadtree-depth quadtree)
+          (quadtree-size quadtree)))
+
+(set-record-type-printer! <quadtree> display-quadtree)
+
+(define* (make-quadtree bounds #:key (max-size 5) (max-depth 4))
+  "Return a new quadtree that covers the area BOUNDS.  Each node will
+try to hold at maximum MAX-SIZE objects and the tree depth will be
+restricted to MAX-DEPTH."
+  (%make-quadtree bounds max-depth max-size 0 0 (make-vector max-size #f) #t))
+
+(define (quadtree-empty? quadtree)
+  "Return #t if QUADTREE has no objects."
+  (= (quadtree-size quadtree) 0))
+
+(define (quadtree-empty-leaf? quadtree)
+  "Return #t if QUADTREE is an empty leaf node."
+  (and (quadtree-empty? quadtree) (quadtree-leaf? quadtree)))
+
+(define (quadtree-full? quadtree)
+  "Return #t if QUADTREE is at or over desired maximum capacity."
+  (>= (quadtree-size quadtree) (quadtree-max-size quadtree)))
+
+(define (quadtree-max-depth? quadtree)
+  "Return #t if QUADTREE is at the maximum allowed depth."
+  (= (quadtree-depth quadtree) (quadtree-max-depth quadtree)))
+
+(define (quadtree-clear! quadtree)
+  "Clear QUADTREE."
+  (vector-fill! (quadtree-objects quadtree) #f)
+  (set-quadtree-size! quadtree 0)
+  (unless (quadtree-leaf? quadtree)
+    (set-quadtree-leaf! quadtree #t)
+    (quadtree-clear! (quadtree-q1 quadtree))
+    (quadtree-clear! (quadtree-q2 quadtree))
+    (quadtree-clear! (quadtree-q3 quadtree))
+    (quadtree-clear! (quadtree-q4 quadtree))))
+
+(define (quadtree-split! quadtree)
+  "Split QUADTREE region into 4 smaller child nodes."
+  (let* ((max-depth (quadtree-max-depth quadtree))
+         (max-size (quadtree-max-size quadtree))
+         (depth (+ (quadtree-depth quadtree) 1))
+         (length (vector-length (quadtree-objects quadtree)))
+         (bounds (quadtree-bounds quadtree))
+         (x (rect-x bounds))
+         (y (rect-y bounds))
+         (hw (/ (rect-width bounds) 2.0))
+         (hh (/ (rect-height bounds) 2.0)))
+    (define (make-node x y)
+      (%make-quadtree (make-rect x y hw hh) max-depth max-size depth 0
+                      (make-vector length #f) #t))
+    (set-quadtree-leaf! quadtree #f)
+    (unless (quadtree-q1 quadtree)
+      (set-quadtree-q1! quadtree (make-node (+ x hw) (+ y hh)))
+      (set-quadtree-q2! quadtree (make-node x (+ y hh)))
+      (set-quadtree-q3! quadtree (make-node x y))
+      (set-quadtree-q4! quadtree (make-node (+ x hw) y)))))
+
+(define (pick-node quadtree rect)
+  "Return the child node of QUADTREE that fully contains RECT, or
+QUADTREE if RECT overlaps multiple child nodes."
+  (let* ((bounds (quadtree-bounds quadtree))
+         (bx (rect-x bounds))
+         (by (rect-y bounds))
+         (bw (rect-width bounds))
+         (bh (rect-height bounds))
+         (mid-x (+ bx (/ bw 2.0)))
+         (mid-y (+ by (/ bh 2.0)))
+         (x (rect-x rect))
+         (y (rect-y rect))
+         (w (rect-width rect))
+         (h (rect-height rect)))
+    (cond
+     ;; Rect does not overlap this node.
+     ((or (>= x (+ bx bw))
+          (<= (+ x w) bx)
+          (>= y (+ by bh))
+          (<= (+ y h) by))
+      #f)
+     ((quadtree-leaf? quadtree)
+      quadtree)
+     ;; Rect is within the left two quadrants: Q2 and Q3.
+     ((and (< x mid-x) (< (+ x w) mid-x))
+      (cond
+       ;; Rect is within Q3.
+       ((and (< y mid-y) (< (+ y h) mid-y))
+        (quadtree-q3 quadtree))
+       ;; Rect is within Q2.
+       ((> y mid-y)
+        (quadtree-q2 quadtree))
+       ;; Rect spans both Q2 and Q3.
+       (else quadtree)))
+     ;; Rect is within the right two quadrants: Q1 and Q4.
+     ((and (> x mid-x))
+      (cond
+       ;; Rect is within Q4.
+       ((and (< y mid-y) (< (+ y h) mid-y))
+        (quadtree-q4 quadtree))
+       ;; Rect is within Q1.
+       ((> y mid-y)
+        (quadtree-q1 quadtree))
+       ;; Rect spans both Q1 and Q4.
+       (else quadtree)))
+     (else quadtree))))
+
+(define (quadtree-add! quadtree rect object)
+  "Add OBJECT to the list of objets in QUADTREE."
+  (let ((objects (quadtree-objects quadtree))
+        (size (quadtree-size quadtree)))
+    (when (= size (vector-length objects))
+      (let* ((new-size (* size 2))
+             (new-objects (make-vector new-size #f)))
+        (let loop ((i 0))
+          (when (< i size)
+            (vector-set! new-objects i (vector-ref objects i))
+            (loop (+ i 1))))
+        (set-quadtree-objects! quadtree new-objects)))
+    (vector-set! (quadtree-objects quadtree) size (cons rect object))
+    (set-quadtree-size! quadtree (+ size 1))))
+
+(define (quadtree-merge-maybe! quadtree)
+  "Remove child nodes if they are all empty leaf nodes."
+  (when (and (quadtree-empty-leaf? (quadtree-q1 quadtree))
+             (quadtree-empty-leaf? (quadtree-q2 quadtree))
+             (quadtree-empty-leaf? (quadtree-q3 quadtree))
+             (quadtree-empty-leaf? (quadtree-q4 quadtree)))
+    ;; We don't actually get rid of the child nodes.  This means that
+    ;; the quadtree can take more memory than it has to, but it also
+    ;; means that the quadtree doesn't allocate needlessly when
+    ;; objects are constantly being added/removed.
+    (set-quadtree-leaf! quadtree #t)))
+
+(define (quadtree-insert! quadtree rect object)
+  "Insert OBJECT with bounding box RECT into QUADTREE."
+  (let ((node (pick-node quadtree rect)))
+    (cond
+     ;; The rect doesn't fit into the parent node.
+     ((not node)
+      #f)
+     ;; The rect fits completely within one of the child nodes, so
+     ;; descend into that node and repeat the process.
+     ((not (eq? quadtree node))
+      (quadtree-insert! node rect object))
+     ;; The node is a leaf node that is at or over the desired
+     ;; capacity, so we need to split it and redistribute the objects.
+     ;; Nodes that have reached the maximum allowed depth cannot be
+     ;; split.
+     ((and (quadtree-full? quadtree)
+           (quadtree-leaf? quadtree)
+           (not (quadtree-max-depth? quadtree)))
+      (let ((objects (quadtree-objects quadtree))
+            (size (quadtree-size quadtree)))
+        (quadtree-split! quadtree)
+        (let loop ((i 0))
+          (if (< i size)
+              (let ((obj (vector-ref objects i)))
+                (loop (+ i 1))
+                (quadtree-insert! quadtree (car obj) (cdr obj)))
+              (begin
+                (set-quadtree-size! quadtree 0)
+                (vector-fill! objects #f))))
+        (quadtree-insert! quadtree rect object)
+        (quadtree-merge-maybe! quadtree)))
+     ;; The node is either under the desired maximum objects threshold
+     ;; or has no choice but to accept another object because there is
+     ;; no child node that fully contains the rect or we have reached
+     ;; the maximum allowed tree depth.
+     (else
+      (quadtree-add! quadtree rect object)))))
+
+(define (quadtree-delete! quadtree rect object)
+  "Delete OBJECT, who occupies the space RECT, from QUADTREE."
+  (let ((node (pick-node quadtree rect)))
+    (cond
+     ((not node)
+      #f)
+     ((eq? quadtree node)
+      (let ((objects (quadtree-objects quadtree))
+            (size (quadtree-size quadtree)))
+        (let loop ((i 0))
+          (cond
+           ((= i size)
+            #f)
+           ((eq? (cdr (vector-ref objects i)) object)
+            (let ((new-size (- size 1)))
+              (vector-set! objects i (vector-ref objects new-size))
+              (vector-set! objects new-size #f)
+              (set-quadtree-size! quadtree new-size)
+              #t))
+           (else
+            (loop (+ i 1)))))))
+     (else
+      (and (quadtree-delete! node rect object)
+           (begin
+             (quadtree-merge-maybe! quadtree)
+             #t))))))
+
+(define (quadtree-find quadtree rect pred)
+  "Return the first object in QUADTREE in the vicinity of RECT that
+satisfies PRED."
+  (let ((objects (quadtree-objects quadtree))
+        (size (quadtree-size quadtree))
+        (next-node (pick-node quadtree rect)))
+    (if next-node
+        (let loop ((i 0))
+          (if (< i size)
+              (let ((object (cdr (vector-ref objects i))))
+                (if (pred object)
+                    object
+                    (loop (+ i 1))))
+              (cond
+               ((and (eq? next-node quadtree) (not (quadtree-leaf? quadtree)))
+                (or (quadtree-find (quadtree-q1 quadtree) rect pred)
+                    (quadtree-find (quadtree-q2 quadtree) rect pred)
+                    (quadtree-find (quadtree-q3 quadtree) rect pred)
+                    (quadtree-find (quadtree-q4 quadtree) rect pred)))
+               ((eq? next-node quadtree)
+                #f)
+               (else
+                (quadtree-find next-node rect pred)))))
+        #f)))
+
+(define (quadtree-fold quadtree rect init proc)
+  "Apply PROC to all objects in the vicinity of RECT in QUADTREE to
+build a result and return that result.  INIT is the initial result.
+If there are no objects in the vicinity of RECT, just INIT is
+returned."
+  (let ((objects (quadtree-objects quadtree))
+        (size (quadtree-size quadtree))
+        (next-node (pick-node quadtree rect)))
+    (if next-node
+        (let loop ((i 0)
+                   (memo init))
+          (if (< i size)
+              (loop (+ i 1) (proc (cdr (vector-ref objects i)) memo))
+              (cond
+               ((and (eq? next-node quadtree) (not (quadtree-leaf? quadtree)))
+                (let* ((q1 (quadtree-fold (quadtree-q1 quadtree) rect memo proc))
+                       (q2 (quadtree-fold (quadtree-q2 quadtree) rect q1 proc))
+                       (q3 (quadtree-fold (quadtree-q3 quadtree) rect q2 proc)))
+                  (quadtree-fold (quadtree-q4 quadtree) rect q3 proc)))
+               ((eq? next-node quadtree)
+                memo)
+               (else
+                (quadtree-fold next-node rect memo proc)))))
+        init)))
+
+(define (quadtree-for-each quadtree rect proc)
+  "Call PROC for all objects in the vicinity of RECT in QUADTREE."
+  (let ((objects (quadtree-objects quadtree))
+        (size (quadtree-size quadtree))
+        (next-node (pick-node quadtree rect)))
+    (when next-node
+      (let loop ((i 0))
+        (when (< i size)
+          (proc (cdr (vector-ref objects i)))
+          (loop (+ i 1))))
+      (cond
+       ((and (eq? next-node quadtree) (not (quadtree-leaf? quadtree)))
+        (quadtree-for-each (quadtree-q1 quadtree) rect proc)
+        (quadtree-for-each (quadtree-q2 quadtree) rect proc)
+        (quadtree-for-each (quadtree-q3 quadtree) rect proc)
+        (quadtree-for-each (quadtree-q4 quadtree) rect proc))
+       ((eq? next-node quadtree)
+        *unspecified*)
+       (else
+        (quadtree-for-each next-node rect proc))))))