Start of LFErlang octree implementation

octree.clj

(defmodule octree
  (export all))
(include-file "all2.lfe") ; lfe_utils library


;; Vectors:

(defmacro :vec3 ((x y z) `(tuple 'vec3 ,x ,y ,z)))

(defn :+ [(:vec3 x y z) (:vec3 a b c)]
  (:vec3 (+ x a) (+ y b) (+ z c)))

(defn null-vec [] (:vec3 0 0 0))


;; Bit helpers:

; Test positive integers.
(defn pow2? 
  [n] (when (=< n #xff)) (fst-8-bits? n)
  [n] (andalso (== 0 (band #xff n)) 
               (pow2? (bsr n 8))))

; Test positive bytes.
(defn fst-8-bits?
  [1] 'true [2] 'true [4] 'true [8] 'true
  [16] 'true [32] 'true [64] 'true [128] 'true
  [_] 'false)

(defn pow2-vec? [(:vec3 x y z)] 
  (andalso (pow2? x) (pow2? y) (pow2? z)))

(defn bsr-vec [(:vec3 x y z) n]
  (:vec3 (bsr x n)
         (bsr y n)
         (bsr z n)))


;; Octree: (well, predetermined BSP tree)

(defmacro :octree ((size data) `(tuple 'octree ,size ,data)))

; Construct octree using an (:vec3 x y z) --> volume-type lookup fun.
(defn octree [size lup] 
  (if (pow2-vec? size) ; dimensions must be powers of two
    (:octree size (build size (null-vec) 'x lup))
    (error 'badarg (list size lup))))
; where
  (defn build 
    ; volumes are 1^3 or larger in size
    [(:vec3 1 1 1) pos _ lup] (funcall lup pos)
    ; input is recursively split in two along X,Y,Z,X,Y,Z,X...
    [size pos axis lup] 
      (in (if (== left right)
            left ; adjacent identical volumes are merged into a larger one
            (cons left right))
        [left (build size' pos axis' lup)
         right (build size' (:+ pos offset) axis' lup)
         (cons size' offset) (split axis size)
         axis' (split-order axis)]))
  ; Split a vector on one axis, returning result and offset of split.
  (defn split
    ['x (:vec3 1 y z)] (cons (:vec3 1 y z) (null-vec))
    ['y (:vec3 x 1 z)] (cons (:vec3 x 1 z) (null-vec))
    ['z (:vec3 x y 1)] (cons (:vec3 x y 1) (null-vec))
    ['x (:vec3 x y z)] (cons (:vec3 (bsr x 1) y z) (:vec3 (bsr x 1) 0 0))
    ['y (:vec3 x y z)] (cons (:vec3 x (bsr y 1) z) (:vec3 0 (bsr y 1) 0))
    ['z (:vec3 x y z)] (cons (:vec3 x y (bsr z 1)) (:vec3 0 0 (bsr z 1))))
  (defn split-order
    ['x] 'y ['y] 'z ['z] 'x)

; Look up a coordinate.
(defn at 
  [pos (:octree size data)] (when (is_list data))
  (let* [((:vec3 x y z) pos)
         (mask (bsr-vec size 1))
         ((:vec3 xm ym zm) mask)]
    (-> data
        (choose-split x xm <>)
        (choose-split y ym <>)
        (choose-split z zm <>)
        (at pos (:octree mask <>))))
  [_ (:octree _ volume)] volume)
; where
  (defn choose-split 
    [path mask data] (when (is_list data))
      (if (== 0 (band mask path)) (car data) (cdr data))
    ; can't split leaves of tree (volumes)
    [_ _ data] data)


; Structure-preserving map.
(defn smap [f (:octree size data)] 
  (:octree size (consmap f data)))
; where
  (defn consmap 
    [f (cons l r)] 
      (in (if (/= l' r')
            (cons l' r')
            l') ; simplify cells if possible
        [l' (consmap f l) 
         r' (consmap f r)])
    [f x] (funcall f x))


;; Test functions:

; S4=[0,1,2,3],[octree:at({vec3,X,Y,Z},octree:test4())||Z<-S4,Y<-S4,X<-S4].
(defn test4 []
  (in (octree size data)
    [size (:vec3 4 4 4)
     data (fn [(:vec3 x y z)] 
            (: binary at bin' (+ x (+ (* 4 y) (* 16 z)))))
     bin' (-> bin binary_to_list (: lists map chr->num <>) list_to_binary)
     chr->num (fn [c] (- c 48))
     bin (binary "2000" "0000" "0000" "0000"
                 "0000" "0000" "0000" "0000"
                 "0000" "0000" "0011" "0011"
                 "0000" "0000" "0011" "0011")]))

; S2=[0,1],[octree:at({vec3,X,Y,Z},octree:test2())||Z<-S2,Y<-S2,X<-S2].
(defn test2 []
  (in (octree size data)
    [size (:vec3 2 2 2)
     data (fn [(:vec3 x y z)] 
            (: binary at bin' (+ x (+ (* 2 y) (* 4 z)))))
     bin' (-> bin binary_to_list (: lists map chr->num <>) list_to_binary)
     chr->num (fn [c] (- c 48))
     bin (binary "12" 
                 "34" "56" 
                      "78")]))
(defn bin_lut [n bin]
  (fn [(:vec3 x y z)]
    (: binary at bin (+ x (+ (* n y) (* (* n n) z))))))


;; Debug:

; Pretty-print tree structure.
(defn print [(:octree _ (cons l r))] 
  (print-rec `(,(cons l r)) 'x))
; where
  (defn print-rec [data axis] 
    ; recurse until no splits left in structure
    (if (: lists any (cut is_list <>) data) 
      (let* [(flatten (fn [(cons l r) acc] (cons l (cons r acc))
                          [x acc] (cons x acc)))
             ; flatten one split per recursion
             (split (: lists foldr flatten '() data))]
        (: io format '"~p split: ~p~n" `(,axis ,split))
        (print-rec split (split-order axis)))
      'ok))