module Resource.Gltf.Scene where
import RIO
import Codec.GlTF qualified as GlTF
import Codec.GlTF.Mesh qualified as GlTF (MeshIx(..))
import Codec.GlTF.Node qualified as GlTF (Node, NodeIx(..))
import Codec.GlTF.Node qualified as Node
import Codec.GlTF.Root qualified as Root
import Codec.GlTF.Scene qualified as Scene
import Data.Coerce (coerce)
import Data.Tree (Tree)
import Data.Tree qualified as Tree
import Geomancy (Transform(..), quaternion)
import Geomancy.Mat4 qualified as Mat4
import Geomancy.Transform qualified as Transform
import Geomancy.Vec3 qualified as Vec3
import RIO.Vector qualified as Vector
import Resource.Gltf.Model (Mesh, MeshPrimitive, Stuff(..), VertexAttrs(..))
unfoldSceneM
:: HasLogFunc env
=> Int
-> Transform
-> Root.GlTF
-> Vector Mesh
-> RIO env (Tree SceneNode)
unfoldSceneM materialOffset initialTransform root allMeshes = do
(allNodes, initialNodes) <- getRootNodes root
meshes <- unfoldNodesM allMeshes allNodes (toList initialNodes)
pure Tree.Node
{ rootLabel = SceneNode
{ snOrigin = 0
, snPrimitives = Nothing
, snNode = emptyNode
}
, subForest =
map (toSceneNode materialOffset initialTransform) meshes
}
{-# INLINEABLE toSceneNode #-}
toSceneNode
:: Int
-> Transform
-> Tree (Maybe (Vector MeshPrimitive), GlTF.Node)
-> Tree SceneNode
toSceneNode materialOffset initialTransform =
fmap (injectTransforms materialOffset) .
collectTransforms initialTransform
getRootNodes :: HasLogFunc env => GlTF.GlTF -> RIO env (Vector Node.Node, Vector Node.NodeIx)
getRootNodes root = do
rootNodes <- case Root.scenes root of
Nothing -> do
logWarn "No scenes"
pure Nothing
Just scenes ->
case Vector.toList scenes of
[] -> do
logWarn "Empty scene vector"
pure Nothing
[one] ->
pure $ Scene.nodes one
pick : _rest -> do
logWarn $ mconcat
[ "Picking first scene among "
, display (Vector.length scenes)
]
pure $ Scene.nodes pick
allNodes <- case Root.nodes root of
Nothing ->
throwString "TODO: fallback to raw meshes"
Just nodes -> do
pure nodes
case rootNodes of
Nothing ->
throwString "TODO: fallback for lack of scene"
Just start -> do
pure (allNodes, start)
emptyNode :: Node.Node
emptyNode = Node.Node
{ camera = Nothing
, children = Nothing
, skin = Nothing
, matrix = Nothing
, mesh = Nothing
, rotation = Nothing
, scale = Nothing
, translation = Nothing
, weights = Nothing
, name = Nothing
, extensions = Nothing
, extras = Nothing
}
data LookupError
= NodeNotFound Int
| MeshNotFound Int
deriving (Eq, Ord, Show)
instance Exception LookupError
injectTransforms
:: Int
-> ( Maybe (Vector MeshPrimitive)
, Transform
, GlTF.Node
)
-> SceneNode
injectTransforms materialOffset (mmesh, transform, snNode) = SceneNode{..}
where
snOrigin = Vec3.Packed $ Transform.apply 0 transform
snPrimitives = fmap (Vector.map $ bimap adjustMaterial adjustNode) mmesh
adjustMaterial = fmap \(materialId, gltfMaterial) ->
( materialId + materialOffset
, gltfMaterial
)
adjustNode Stuff{..} = Stuff
{ sPositions = fmap applyTransform sPositions
, sAttrs = fmap applyAttrTransform sAttrs
, sIndices = sIndices
}
where
applyTransform pos =
coerce $ Transform.apply (coerce pos) transform
applyTransformDir dir =
coerce $ Vec3.normalize $ Transform.apply (coerce dir) (Transform transformDir)
where
transformDir = Mat4.transpose . Mat4.inverse $
Mat4.pointwise (unTransform transform) nullifyTranslation (*)
applyAttrTransform va = va
{ vaNormal = applyTransformDir (vaNormal va)
, vaTangent = applyTransformDir (vaTangent va)
}
collectTransforms
:: Transform
-> Tree (Maybe (Vector MeshPrimitive), GlTF.Node)
-> Tree (Maybe (Vector MeshPrimitive), Transform, GlTF.Node)
collectTransforms initial root = Tree.unfoldTree go (initial, root)
where
go (parent, Tree.Node{rootLabel=(mmesh, node), subForest}) =
let
collected = localTransform node <> parent
in
( (mmesh, collected, node)
, map (collected,) subForest
)
-- | Build node tree and shed lookup errors as exception.
unfoldNodesM
:: MonadThrow m
=> Vector Mesh
-> Vector GlTF.Node
-> [GlTF.NodeIx]
-> m [Tree (Maybe (Vector MeshPrimitive), GlTF.Node)]
unfoldNodesM allMeshes allNodes =
either throwM pure . traverse (inflateNode allMeshes allNodes)
-- | Combine lookup operations for nodes and meshes.
inflateNode
:: Vector Mesh
-> Vector GlTF.Node
-> GlTF.NodeIx
-> Either LookupError (Tree (Maybe (Vector MeshPrimitive), GlTF.Node))
inflateNode allMeshes allNodes startNode =
for (unfoldNode allNodes startNode) \getNode -> do
node <- getNode
mesh <- getMesh allMeshes node
pure (mesh, node)
unfoldNode
:: Vector GlTF.Node
-> GlTF.NodeIx
-> Tree (Either LookupError GlTF.Node)
unfoldNode allNodes = Tree.unfoldTree fetch
where
fetch (GlTF.NodeIx ix) =
case allNodes Vector.!? ix of
Nothing ->
( Left $ NodeNotFound ix
, []
)
Just node ->
( Right node
, maybe [] toList $ Node.children node
)
getMesh
:: Vector mesh
-> GlTF.Node
-> Either LookupError (Maybe mesh)
getMesh allMeshes node =
case Node.mesh node of
Nothing ->
Right Nothing
Just (GlTF.MeshIx ix) ->
case allMeshes Vector.!? ix of
Nothing ->
Left $ MeshNotFound ix
Just mesh ->
Right (Just mesh)
localTransform :: Node.Node -> Transform
localTransform node = mconcat
[ nodeMatrix
, nodeScale, nodeRotate, nodeTranslate
]
where
nodeScale = case Node.scale node of
Nothing ->
mempty
Just (sx, sy, sz) ->
Transform.scale3 sx sy sz
nodeTranslate =
case Node.translation node of
Nothing ->
mempty
Just (tx, ty, tz) ->
Transform.translate tx ty tz
nodeRotate = case Node.rotation node of
Nothing ->
mempty
Just (qx, qy, qz, qW) ->
Transform.rotateQ $ quaternion qW qx qy qz
nodeMatrix = case Node.matrix node of
Nothing ->
mempty
Just (Node.NodeMatrix memoryBytes) ->
case Vector.toList memoryBytes of
[ x1, x2, x3, x4,
y1, y2, y3, y4,
z1, z2, z3, z4,
w1, w2, w3, w4 ] ->
Mat4.colMajor @Transform
x1 y1 z1 w1
x2 y2 z2 w2
x3 y3 z3 w3
x4 y4 z4 w4
_ ->
error "Node matrix isn't 16-element"
nullifyTranslation :: Mat4.Mat4
nullifyTranslation =
Mat4.rowMajor
1 1 1 0
1 1 1 0
1 1 1 0
0 0 0 1
data SceneNode = SceneNode
{ snOrigin :: ~Vec3.Packed
, snPrimitives :: Maybe (Vector MeshPrimitive)
, snNode :: GlTF.Node
}
deriving (Show)