Octree.hs

{-# language NamedFieldPuns #-}
{-# language GeneralizedNewtypeDeriving #-}

import           Control.Monad                 (when)
import           Control.Monad.IO.Class        (liftIO)
import           Control.Monad.Reader          (ReaderT(..), asks)
import           Control.Monad.Primitive       (PrimMonad, PrimState)

import           Data.Bits                     (testBit, (.|.), shiftL)
import           Data.List                     (intercalate)
import           Data.Traversable              (for)
import qualified Data.Vector.Mutable      as V

-- | Number of the octant.
newtype Dir = Dir Int
    deriving (Eq, Ord, Enum, Show, Num, Real, Integral)

-- | Predicates to check of we select farther octant in given direction.
isX, isY, isZ :: Dir -> Bool
isX (Dir dir) = testBit dir 0
isY (Dir dir) = testBit dir 1
isZ (Dir dir) = testBit dir 2

-- | Point in space.
data Point = Point { x, y, z :: Int }
    deriving (Eq)
    
instance Show Point where
    show Point {x, y, z} = "{" ++ show x ++ "," ++ show y ++ "," ++ show z ++ "}"

-- | Level-of-detail we're on.
newtype Level = Level Int
    deriving (Eq, Ord, Enum, Num, Real, Integral)

instance Show Level where
    show (Level level) = show level

-- | Get imminent octant to go in.
octantAt :: Level -> Point -> Dir
octantAt (Level level) Point {x, y, z} =
    Dir ( (if testBit x level then 1 else 0)
        + (if testBit y level then 2 else 0)
        + (if testBit z level then 4 else 0)
        )
    
-- | Transform the origin of the 'Location' to one of given octant.
descentPoint :: Dir -> Level -> Point -> Point
descentPoint (Dir dir) (Level level) Point {x, y, z} =
    Point
        { x = x .|. if testBit dir 0 then shiftL 1 level else 0
        , y = y .|. if testBit dir 1 then shiftL 1 level else 0
        , z = z .|. if testBit dir 2 then shiftL 1 level else 0
        }

-- | Cubic volume in space. The 'level' is a logarithm of edge length.
data Location = Location { point :: Point, level :: Level }
    deriving (Eq)
    
instance Show Location where
    show Location {point, level} = show point ++ "@" ++ show (2 ^ level)

-- | Find octant containing given point inside current location.
octantTo :: Point -> Location -> Dir
octantTo point Location { level } = octantAt level point

-- | Transform location to one of the given octant.
descentLocation :: Dir -> Location -> Location
descentLocation dir Location { point, level } =
    let level' = level - 1
    in  Location
        { point = descentPoint dir level point
        , level = level'
        }

-- | Are we just above finest of atoms?
isAtomic :: Location -> Bool
isAtomic Location { level } = level == 0

-- | Space separation.
data Octree a
    -- | Indivisible space (cubical chunk of air, stone, etc)
    = Atom a
    
    -- | Divisible space: 8 different things
    | Split Location (V.IOVector (Octree a))
    
    -- | Space, that was not generated yet
    | Thunk Location

isAtom :: Octree a -> Bool
isAtom (Atom {}) = True
isAtom  _        = False

-- | Context for operations over 'Octree'.
data OctreeCtx a = OctreeCtx
    { gen :: Location -> M a (Octree a)
    }

-- | IO enriched with readonly access to 'OctreeCtx'.
type M a = ReaderT (OctreeCtx a) IO

-- | Convenent name for spooky UNSAFE reading from array.
infix 1 ?
(?) :: (Integral i, PrimMonad m) => V.MVector (PrimState m) a -> i -> m a
vec ? index = V.unsafeRead vec (fromIntegral index)

-- | Convenent name for spooky UNSAFE writing into array.
--
--   Use as: array '=:' index '$' value
infix 1 =:
(=:) :: (Integral i, PrimMonad m) => V.MVector (PrimState m) a -> i -> a -> m ()
(vec =: index) value = V.unsafeWrite vec (fromIntegral index) value

-- | Descent into given octant of the 'Octree'.
descent :: Dir -> Octree a -> M a (Octree a)
descent (Dir dir) tree = do
--    put $ "descent " ++ show dir

    case tree of
        Split loc children -> do
            child <- children ? dir

            case child of
                Thunk pos -> do
                    generate <- asks gen
                    child'   <- generate pos
                    children =: dir $ child'
                    return child'

                other -> do
                    return other
        
        other -> 
            error "Octree.descent: required Split"

-- | Get material at given point.
at :: Point -> Octree a -> M a a
at point tree = do
--    put $ "at " ++ show point
    case tree of
        Split loc children -> do
            let dir = octantTo point loc
            child <- descent dir tree
            at point child
            
        Atom a -> do
            return a
            
        Thunk gen -> do
            error "Octree.at: Thunk"

-- | Set material at given point.
assign :: Eq a => Point -> a -> Octree a -> M a ()
assign point material tree = case tree of
    Split loc children -> do
        let dir = octantTo point loc
        
        if isAtomic loc
        then do
            children =: dir $ Atom material

        else do
            child  <- descent dir tree
            child' <- case child of
                Atom a -> do
                    grandchildren <- V.new 8
                    V.set grandchildren (Atom a)
                    let split = Split (descentLocation dir loc) grandchildren
                    children =: dir $ split
                    return split
                
                other -> do
                    return other

            assign point material child'
            
            tryWelding dir children
            
    other ->
        error "Octree.assign: required Split"
            
-- | Optimise after write: if there are 8 equal atoms, merge into one.
tryWelding :: Eq a => Dir -> V.IOVector (Octree a) -> M a ()
tryWelding dir children = do
    subtree <- children ? dir
    case subtree of
        Split loc children -> do
            yes <- flip allM [0.. 7] $ \i -> do
                child <- children ? i
                return (isAtom child)
                
            when yes $ do
                mat : rest <- for [0.. 7] $ \i -> do
                    Atom mat <- children ? i
                    return mat
                
                when (all (== mat) rest) $ do
                    children =: dir $ Atom mat

        other -> do
            return ()

-- | Why isn't this in @base@ package?            
allM :: (Monad m, Traversable f) => (a -> m Bool) -> f a -> m Bool
allM pred container = do
    -- TODO: implement properly.
    bools <- mapM pred container
    return (and bools)

generate :: Location -> M Int (Octree Int)
generate loc = do
    children <- V.new 8
    if isAtomic loc
    then do
        for [0.. 7] $ \i -> do
            children =: i $ Atom i
        return $ Split loc children
    else do
        for [0.. 7] $ \i -> do
            children =: i $ Thunk $ descentLocation (Dir i) loc
            
        return (Split loc children)

-- | Create tree at given location using generator from r/o ctx.
makeAt :: Location -> M a (Octree a)
makeAt loc = do
    generate <- asks gen
    generate loc

-- | ToString.
dump :: Show a => Octree a -> M a String
dump tree = case tree of
    Atom a -> do
        return (show a)

    Split loc children -> do
        strs <- for [0.. 7] $ \i -> do
            child <- children ? i
            dump child
            
        return (show loc ++ ":[" ++ intercalate "," strs ++ "]")

    Thunk _ -> do
        return "?"

put :: String -> M a ()
put s = liftIO (putStrLn s)

printTree :: Show a => Octree a -> M a ()
printTree tree = do
    str <- dump tree
    put str

main = do
    flip runReaderT (OctreeCtx generate) $ do
        tree <- makeAt $ Location (Point 0 0 0) 3
        _    <- printTree tree
        a    <- at (Point 0 0 0) tree
        _    <- put (show a)
        _    <- printTree tree
        b    <- at (Point 1 0 1) tree
        _    <- put (show b)
        _    <- printTree tree
        b    <- at (Point (1 + 128) (0 + 16) (1 + 64)) tree
        _    <- put (show b)
        _    <- printTree tree
        _    <- assign (Point 5 3 2) 42 tree
        _    <- printTree tree
        b    <- at (Point 5 3 2) tree
        _    <- put (show b)
        _    <- printTree tree
        return ()