Chunk.hs

{- | A chunk of 16x16x16 atoms.
-}

module Chunk where

import Prelude hiding (read)

import Control.Monad
import Control.Monad.Primitive

import Data.IORef
import Data.Functor
import Data.Foldable
import Data.Traversable

import Data.Vector.Unboxed (Unbox)
import qualified Data.Vector.Grow.Unboxed as Vector
import Data.Vector.Grow.Unboxed (GrowVector)

class PrimMonad m => HasVar m where
  type Var m :: * -> *
  alloc :: a -> m (Var m a)
  fetch :: Var m a -> m a
  ($=)  :: Var m a -> a -> m ()

instance HasVar IO where
  type Var IO = IORef
  alloc = newIORef
  fetch = readIORef
  ($=)  = writeIORef

{- | A chunk is either a vector with size 1 (then it is a monolit),
     or a vector with size 16^3 (then it is normal chunk).

     Any write will turn monolit into normal chunk.
-}
data Chunk m a = Chunk { mono :: Var m Bool, rawChunk :: GrowVector (PrimState m) a }

{- | Will do for now.
-}
type Stored = Unbox

{- | Create a monolitic chunk.
-}
monolit :: (HasVar m, Stored a) => a -> m (Chunk m a)
monolit a = do
  mono <- alloc True
  gv   <- Vector.newSized 1 1
  Vector.unsafeWrite gv 0 a
  return $ Chunk mono gv

{- | Create a normal chunk.

     The count of elements provided must be 16^3 or more.

     Excess elements will be ignored.
-}
chunk :: (HasVar m, Stored a) => [a] -> m (Chunk m a)
chunk as = do
  mono <- alloc False
  gv   <- Vector.newSized 0 (16 * 16 * 16)
  unsafeFill (Chunk mono gv) (16 * 16 * 16) as
  return $ Chunk mono gv
  where

{- | Create a normal chunk.

     The count of elements provided must be 16^3 or more.

     Excess elements will be ignored.
-}
unsafeFill :: (HasVar m, Stored a) => Chunk m a -> Int -> [a] -> m ()
unsafeFill (Chunk _ gv) = fill
  where
    fill n (a : as) | n > 0 = do
      Vector.unsafePushBack gv a
      fill (n - 1) as
    fill _ _ = return ()

{- | Check if chunk is a monolit.
-}
isMonolit :: (HasVar m, Stored a) => Chunk m a -> m Bool
isMonolit (Chunk mono _) = do
  fetch mono

{- | Get element of a chunk.
-}
{-# INLINE read #-}
read :: (HasVar m, Stored a) => Chunk m a -> Int -> m a
read c@(Chunk _ raw) i = do
  mono <- isMonolit c
  Vector.unsafeRead raw if mono then 0 else i

{- | Get many elements of a chunk.
-}
{-# INLINE massRead #-}
massRead :: (HasVar m, Stored a) => Chunk m a -> [Int] -> m [a]
massRead c@(Chunk _ raw) is = do
  mono <- isMonolit c
  if mono
  then do
    a <- Vector.unsafeRead raw 0
    return $ a <$ is

  else do
    for is $ Vector.unsafeRead raw

{- | Write element into of a chunk.

     Will transform monolit chunk into normal one.
-}
{-# INLINE write #-}
write :: (HasVar m, Stored a) => Chunk m a -> Int -> a -> m ()
write c@(Chunk _ raw) i a = do
  mono <- isMonolit c
  when mono do
    unsafeToChunk c

  Vector.write raw i a

unsafeToChunk :: (HasVar m, Stored a) => Chunk m a -> m ()
unsafeToChunk c@(Chunk mono raw) = do
  old <- read c 0
  Vector.ensure raw (16 * 16 * 16)
  unsafeFill c (16 * 16 * 16 - 1) (repeat old)
  mono $= False

{- | Perform a batch-write element into of a chunk.

     Will check if chunk is a monolit only once.
-}
{-# INLINE massWrite #-}
massWrite :: (HasVar m, Stored a) => Chunk m a -> [(Int, a)] -> m ()
massWrite c@(Chunk _ raw) batch = do
  mono <- isMonolit c
  when mono do
    unsafeToChunk c

  for_ batch \(i, a) -> do
    Vector.unsafeWrite raw i a

Octree.hs

{- | A sparse voxel octree.

     The tree carries some monoidal fold in each node.
     Use it to store average node color or what not.

     It has /neither/ the notion of its own size,
     /nor/ the notion of maximal granularity,
     /nor/ the coordinates.
     You have to watch them yourself, probably by keeping
     the current log8 of volume covered by tree.

     If you load up a segment and it is outside a tree,
     expand in general direction of a segment then put it in.

     If you want to remove a segment from memory, replace it
     with some atom of your choice.

     The path used to access is a @Point x y z@ but
     it is @(x y z)@ matrix, bit-transposed.

     The tree will self-optimise (and there is an off-flag).
     If the updated branch contains atoms with identical contents,
     it will be replaced with one atom.
-}

module Octree
  ( -- * Octree type and constructors
    Octree (..)
  , atom
  , branch
  , expand

    -- * High-level access
  , Path
  , location
  , batched
  , scanned

    -- * Low-level access
  , ZipperT
  , runZipperT
  , exit
  , step
  , go
  , change
  , here
  , batch
  , scan
  ) where

import Control.Lens
import Control.Monad.Extra
import Control.Monad.State
import Data.Foldable
import qualified Data.Vector as Vector
import Data.Vector ((!), (//), Vector)
import GHC.Generics

{- | An octree of atoms @a@ and some previev (color, material hp) @b@.
-}
data Octree b a
  = Branch !b (Vector (Octree b a))
  | Atom   !b a
  deriving stock (Show, Generic)

makePrisms ''Octree

{- | Regenerate or extract previews. Happens automatically on updates.
-}
class Monoid b => Draft a b | a -> b where
  draft :: a -> b

instance Draft a b => Draft (Octree b a) b where
  draft = \case
    Branch b _ -> b
    Atom   b _ -> b

{- | Peano-cube merge of vector.
-}
merge8 :: Monoid a => Vector a -> a
merge8 v = ((v ! 0 <> v ! 1) <> (v ! 2 <> v ! 3)) <> ((v ! 4 <> v ! 5) <> (v ! 6 <> v ! 7))

{- | Each `Int` is 3-bits of @(x, y, z)@ of the same position, starting from high bits.

     Do /not/ use numbers outside @[0.. 7]@, you have been warned.
-}
type Path = [Int]

{- | TODO: make it a CPP-pragma.
-}
selfCompact :: Bool
selfCompact = True

{- | Access lens.
-}
location :: forall a b. (Eq a, Draft a b) => Path -> Lens' (Octree b a) (Octree b a)
location path = lens (getA path) (setA path)
  where
    getA :: Path -> Octree b a -> Octree b a
    getA p s = case (p, s) of
      (i : p', Branch _ v) -> getA p' (v ! i)
      _                    -> s


    setA :: Path -> Octree b a -> Octree b a -> Octree b a
    setA p s a = case (p, s) of
      (i : p', Branch _ v) -> do
        let e  = v ! i
        let e' = setA p' e' a
        let v' = v // [(i, e')]
        if selfCompact
        then
          case allSameAtoms v' of
            Just e' -> atom e'
            _       -> branch v'
        else
          branch v'

      (i : p', Atom _ e) ->
        setA p (split e) a

      ([], _) -> a

    allSameAtoms :: Vector (Octree b a) -> Maybe a
    allSameAtoms v = do
      v' <- traverse (^?_Atom._2) v
      guard $ Prelude.all (\i -> (v' ! 0) == (v' ! i)) [1.. 7]
      return $ v' ! 0

    split :: a -> Octree b a
    split = branch . Vector.replicate 8 . atom

{- | Create one atom.
-}
atom :: Draft a b => a -> Octree b a
atom a = Atom (draft a) a

{- | Create a branch.
-}
branch :: Draft a b => Vector (Octree b a) -> Octree b a
branch es = Branch (merge8 $ fmap draft es) es

{- | Add some space around the octree.
-}
expand :: Draft a b => a -> Int -> Octree b a -> Octree b a
expand vacuum octant octree =
  branch $ Vector.generate 8 $ \i ->
    if i == octant
    then octree
    else atom vacuum

data ZipLayer b a = ZipLayer
  { _tree  :: Octree b a
  , _back  :: Int
  , _dirty :: Bool
  }

makeLenses ''ZipLayer

{- | Run batched sequence of updates.
-}
batched :: (Eq a, Draft a b, MonadFail m) => [(Path, Octree b a -> m (Octree b a))] -> Octree b a -> m (Octree b a)
batched action octree = runZipperT octree (batch action >> exit)

{- | Run a batched sequence of lookups.
-}
scanned :: (Eq a, Draft a b, MonadFail m, Monoid r) => [(Path, Octree b a -> m r)] -> Octree b a -> m r
scanned action octree = runZipperT octree (scan action)

{- | An iterator over the tree.
-}
type ZipperT b a = StateT [ZipLayer b a]

{- | Run the iterator.
-}
runZipperT :: Monad m => Octree b a -> ZipperT b a m x -> m x
runZipperT octree action = evalStateT action $ enter octree

{- | Open the tree.
-}
enter :: Octree b a -> [ZipLayer b a]
enter _tree = [ZipLayer { _tree, _back = 0, _dirty = False }]

{- | Reconstruct current tree.
-}
exit :: (Eq a, Draft a b, MonadFail m) => ZipperT b a m (Octree b a)
exit = do
  exit'
  Just locus <- gets (^?_head.tree)
  return locus
  where
    exit' = do
      gets (^?_tail._head) >>= \case
        Nothing -> return ()
        _       -> do
          step Nothing
          exit'

{- | Perform step into specified direction. On `Nothing` step up.
-}
step :: (Eq a, Draft a b, Monad m) => Maybe Int -> ZipperT b a m ()
step (Just i) = do
  get >>= \case
    zl : _ ->
      modify $ (:) $ ZipLayer
        { _tree  = zl^.tree.location [i]
        , _back  = i
        , _dirty = False
        }
step _ = do
  get >>= \case
    zl : rest -> do
      modify tail
      when (zl^.dirty) do
        modify $ _head
          %~ (dirty .~ True)
          .  (tree.location [zl^.back] .~ zl^.tree)

{- | Perform given sequence of steps.
-}
go :: (Eq a, Draft a b, Monad m) => [Maybe Int] -> ZipperT b a m ()
go = traverse_ step

{- | Register a change.
-}
change :: MonadFail m => (Octree b a -> m (Octree b a)) -> ZipperT b a m ()
change act = do
  Just locus <- gets (^?_head.tree)
  locus' <- lift $ act locus
  modify $ _head %~ ((tree .~ locus') . (dirty .~ True))

{- | Get current subtree.
-}
here :: MonadFail m => ZipperT b a m (Octree b a)
here = do
  Just locus <- gets (^?_head.tree)
  return locus

{- | Run several actions in a batch /from current subtree/. Ends somewhere on the last path.
-}
batch :: (Eq a, Draft a b, MonadFail m) => [(Path, Octree b a -> m (Octree b a))] -> ZipperT b a m ()
batch (differentiate -> items) = do
  for_ items \(path, action) -> do
    go path
    change action

{- | Run several scans in a batch /from current subtree/. Ends somewhere on the last path.
-}
scan :: (Eq a, Draft a b, MonadFail m, Monoid r) => [(Path, Octree b a -> m r)] -> ZipperT b a m r
scan (differentiate -> items) = do
  flip mconcatMapM items \(path, action) -> do
    go path
    here >>= lift . action

{- | Optimise path traversal.
-}
differentiate :: [(Path, x)] -> [([Maybe Int], x)]
differentiate items = do
  let (paths, actions) = unzip items
  let paths' = differential paths
  zip paths' actions

differential :: [Path] -> [[Maybe Int]]
differential []       = []
differential (p : ps) = [map Just p] ++ go p ps
  where
    go :: Path -> [Path] -> [[Maybe Int]]
    go p [] = []
    go p (q : ps) =
      (replicate d Nothing ++ map Just q') : go q ps
      where
        (d, q') = delta p q

    delta :: Path -> Path -> (Int, Path)
    delta  []      []              = (0, [])
    delta  p       []              = (length p, [])
    delta  []      q               = (0, q)
    delta (a : p) (b : q) | a == b = delta p q
    delta  p       q               = (length p, q)

package.yaml

dependencies:
- base
- extra
- grow-vector
- lens
- mtl
- primitive
- vector

default-extensions:
- BlockArguments
- ConstraintKinds
- DeriveGeneric
- DerivingStrategies
- FlexibleInstances
- FunctionalDependencies
- LambdaCase
- MultiParamTypeClasses
- NamedFieldPuns
- RankNTypes
- ScopedTypeVariables
- TemplateHaskell
- TypeApplications
- TypeFamilies
- ViewPatterns

library:
  source-dirs:
  - .