packages feed

keid-resource-gltf-0.1.1.0: src/Resource/Gltf/Scene.hs

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.Mat4.Internal qualified as Mat4 (fromMemory)
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   = adjustIndices sIndices
      }
      where
        {-
          https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html#instantiation

          When a mesh primitive uses any triangle-based topology (i.e., triangles, triangle strip, or triangle fan),
          the determinant of the node’s global transform defines the winding order of that primitive.

          If the determinant is a positive value, the winding order triangle faces is counterclockwise;
          in the opposite case, the winding order is clockwise.
        -}
        adjustIndices =
          if Mat4.det transform < 0 then
            Vector.reverse
          else
            id

        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 = mconcat $ parent : nodeTransforms node
      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)

nodeTransforms :: Node.Node -> [Transform]
nodeTransforms node =
  case nodeMatrix of
    -- glTF spec says matrix should not be mixed with transforms...
    Just takesPrecedence ->
      -- ...but when that happens matrix takes precedence.
      [takesPrecedence]
    Nothing ->
      catMaybes
        [ nodeTranslate
        , nodeRotate
        , nodeScale
        ]
  where
    nodeScale =
      node.scale <&> \(sx, sy, sz) ->
        Transform.scale3 sx sy sz

    nodeTranslate =
      node.translation <&> \(tx, ty, tz) ->
          Transform.translate tx ty tz

    nodeRotate =
      node.rotation <&> \(qx, qy, qz, qW) ->
        Transform.rotateQ $ quaternion qW qx qy qz

    nodeMatrix =
      node.matrix <&> \(Node.NodeMatrix memoryBytes) ->
        Transform $ case Vector.toList memoryBytes of
          [ e0, e1, e2, e3,
            e4, e5, e6, e7,
            e8, e9, eA, eB,
            eC, eD, eE, eF ] ->
              Mat4.fromMemory
                e0 e1 e2 e3
                e4 e5 e6 e7
                e8 e9 eA eB
                eC eD eE eF
          _ ->
            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)