path: root/src/Wavefront.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes       #-}
module Wavefront where

import Wavefront.Types
import Wavefront.Lex

import Control.Arrow
import Control.Monad.State
import qualified Data.ByteString.Lazy.Char8 as L
import qualified Data.DList                 as DList
         ;import Data.DList                 (DList)
import Data.Functor.Identity
import qualified Data.IntMap                as IntMap
import qualified Data.Text                  as T
import Data.Text.Encoding                   (decodeUtf8)
import qualified Data.Vector                as Vector
         ;import Data.Vector                (Vector)
import qualified Rank2

type WavefrontOBJ = OBJ Vector

newtype Count x = Count Int


incrementCount :: MonadState s m => (s -> Count x) -> (Count x -> s -> s) -> m ()
incrementCount field setField = do
    Count c0 <- gets field
    let c = succ c0
    c `seq` modify (setField $ Count c)

fixupRef :: Count x -> Int -> Int
fixupRef (Count n) x | x >  0    = x - 1 -- Renumber from 0.
                     | otherwise = n + x -- Negative values are relative.

fixupTriple :: OBJ Count -> RefTriple -> RefTriple
fixupTriple o (RefTriple v t n) =
    RefTriple (fixupRef (objLocations o) v)
              (fixupRef (objTexCoords o) <$> t)
              (fixupRef (objNormals   o) <$> n)

objBookKeeping :: Monad m => ObjBuilder m -> ObjBuilder (StateT (OBJ Count) m)
objBookKeeping builder = (lift Rank2.<$> builder)
    { vertex = \xs -> do lift $ vertex builder xs
                         incrementCount objLocations $ \x o -> o { objLocations = x }
    , vertexT = \xs -> do lift $ vertexT builder xs
                          incrementCount objTexCoords $ \x o -> o { objTexCoords = x }
    , vertexN = \xs -> do lift $ vertexN builder xs
                          incrementCount objNormals $ \x o -> o { objNormals = x }
    , points = \xs -> do
        n <- gets objLocations
        lift $ points builder $ fixupRef n <$> xs
        incrementCount objPoints $ \x o -> o { objPoints = x }
    , line = \ts -> do
        o <- get
        lift $ line builder $ fixupTriple o <$> ts
        incrementCount objLines $ \x o -> o { objLines = x }
    , face = \ts -> do
        o <- get
        lift $ face builder $ fixupTriple o <$> ts
        incrementCount objFaces $ \x o -> o { objFaces = x }
    }

-- TODO: meshlab vertex colors extension
mkv :: [Double] -> Location
mkv [x,y,z,r,g,b] = Location x' y' z' 1 where (x':y':z':_) = map realToFrac [x,y,z]
mkv cs = Location x y z w where (x:y:z:w:_) = map realToFrac cs ++ repeat 1

mkt :: [Double] -> TexCoord
mkt cs = TexCoord x y z where (x:y:z:_) = map realToFrac cs ++ repeat 0

mkn :: [Double] -> Normal
mkn cs = Normal x y z where (x:y:z:_) = map realToFrac cs ++ repeat 0

mkl :: RefTriple -> RefTriple -> Line
mkl (RefTriple a at _) (RefTriple b bt _) = Line (LineIndex a at) (LineIndex b bt)

-- I'd have thought these would be Coercible, but I guess not.
mkF :: RefTriple -> FaceIndex
mkF (RefTriple a at an) = FaceIndex a at an

elemental :: Element () -> x -> Element x
elemental element x = fmap (const x) element

modifyFirst :: MonadState ParserState m => (OBJ DList -> OBJ DList) -> m ()
modifyFirst = modify' . (\f s -> s { pstObj = f (pstObj s) })

modifySecond :: MonadState ParserState m => (Element () -> Element ()) -> m ()
modifySecond = modify' . (\f s -> s { pstElm = f (pstElm s) })

(*.*) :: (OBJ DList -> a) -> (Element () -> b) -> ParserState -> (a, b)
fld *.* elm = \s -> (fld (pstObj s), elm (pstElm s))

data FFPts = Curv Float Float [Int]
           | Curv2 [Int]
           | Surf Float Float Float Float [RefTriple]

data FreeForm = FreeForm
    { ffRat :: Bool
    , ffTyp :: CSType
    , ffDeg :: (Int,Int)
    , ffPts :: FFPts
    }

initFF :: FreeForm
initFF = FreeForm
    { ffRat = False
    , ffTyp = Bspline
    , ffDeg = (1,1)
    , ffPts = Curv2 []
    }

mkcurv2 :: FreeForm -> EmbeddedCurve
mkcurv2 ff = EmbeddedCurve

mkcurv :: FreeForm -> Curve
mkcurv ff = Curve

mksurf :: FreeForm -> Surface
mksurf ff = Surface

addFreeForm :: (forall x. x -> Element x) -> FreeForm -> OBJ DList -> OBJ DList
addFreeForm elm ff o = case ffPts ff of
    Curv2 {} -> o { objEmbeddedCurves = objEmbeddedCurves o `DList.snoc` mkcurv2 ff }
    Curv {}  -> o { objCurves = objCurves o `DList.snoc` elm (mkcurv ff) }
    Surf {}  -> o { objSurfaces = objSurfaces o `DList.snoc` elm (mksurf ff) }

data ParserState = ParserState
    { pstObj :: OBJ DList
    , pstElm :: Element ()
    , pstFF  :: FreeForm
    }

buildOBJ :: ObjBuilder (State ParserState)
buildOBJ = nullBuilder
    { vertex = \xs -> modifyFirst $ \o -> o { objLocations = objLocations o `DList.snoc` mkv xs }
    , vertexT = \xs -> modifyFirst $ \o -> o { objTexCoords = objTexCoords o `DList.snoc` mkt xs }
    , vertexN = \xs -> modifyFirst $ \o -> o { objNormals = objNormals o `DList.snoc` mkn xs }
    , points = \xs -> do
        let p = map Point xs :: [Point]
        (pts,element) <- gets (objPoints *.* elemental)
        modifyFirst $ \o -> o { objPoints = pts `DList.append` fmap element (DList.fromList p) }
    , line = \xs -> do
        (lns,element) <- gets (objLines *.* elemental)
        let l = zipWith mkl xs (tail xs)
        -- Line requires at least two points.  We'll ignore it otherwise.
        when (not $ null l) $
            modifyFirst $ \o -> o { objLines = lns `DList.append` fmap element (DList.fromList l) }
    , face = \xs -> do
        (fcs,element) <- gets (objFaces *.* elemental)
        case map mkF xs of
          a:b:c:ds -> modifyFirst $ \o -> o { objFaces = fcs `DList.snoc` element (Face a b c ds) }
          _        -> return () -- Ignore faces with less than 3 indices.
    , mtllib = \xs -> do
        let l = map decodeUtf8 xs
        libs <- gets (objMtlLibs . pstObj)
        modifyFirst $ \o -> o { objMtlLibs = DList.fromList l `DList.append` libs }
    , groups = \xs -> do
        let g = map decodeUtf8 xs
        modifySecond $ \e -> e { elGroups = g }
    , objectName = \x -> do
        let o = decodeUtf8 x
        modifySecond $ \e -> e { elObject = Just o }
    , usemtl = \x -> do
        let mtl = decodeUtf8 x
        libs <- DList.toList <$> gets (objMtlLibs . pstObj)
        modifySecond $ \e -> e { elMtl = Just (length libs,mtl) }
    , smoothingGroup = \x -> when (x > 0) $ do
        modifySecond $ \e -> e { elSmoothingGroup = fromIntegral x }
    , endFreeForm =
        modify' $ \s -> s { pstObj = addFreeForm (elemental $ pstElm s) (pstFF s) (pstObj s)
                          , pstFF  = initFF }
    }

blankElement :: Element ()
blankElement = Element
            { elObject         = Nothing
            , elGroups         = []
            , elMtl            = Nothing
            , elSmoothingGroup = 0
            , elValue          = ()
            }

emptyCounts :: OBJ Count
emptyCounts = OBJ
            { objLocations = Count 0
            , objTexCoords = Count 0
            , objNormals   = Count 0
            , objEmbeddedCurves = Count 0
            , objPoints    = Count 0
            , objLines     = Count 0
            , objFaces     = Count 0
            , objCurves    = Count 0
            , objSurfaces  = Count 0
            , objMtlLibs   = Count 0
            }

mzeroOBJ :: MonadPlus m => OBJ m
mzeroOBJ = OBJ
            { objLocations = mzero
            , objTexCoords = mzero
            , objNormals   = mzero
            , objEmbeddedCurves = mzero
            , objPoints    = mzero
            , objLines     = mzero
            , objFaces     = mzero
            , objCurves    = mzero
            , objSurfaces  = mzero
            , objMtlLibs   = mzero
            }

parse :: L.ByteString -> OBJ Vector
parse bs = Rank2.fmap (Vector.fromList . DList.toList) obj
 where
    go = parseCustom buildOBJ (const $ return ()) bs
    ParserState { pstObj = obj } = execState go (ParserState mzeroOBJ blankElement initFF)

parseCustom :: Monad m => ObjBuilder m -> (OBJ Count -> m b) -> L.ByteString -> m b
parseCustom builder finish bs = do
    counts <- execStateT (parseOBJ (objBookKeeping builder) (ObjConfig IntMap.empty) bs) emptyCounts
    finish counts

data Renumbering = Renumbering
    { renumV  :: Int -> Int
    , renumVT :: Int -> Int
    , renumVN :: Int -> Int
    , renumVP :: Int -> Int
    }

renumFrom1 :: Renumbering
renumFrom1 = Renumbering
    { renumV  = succ
    , renumVT = succ
    , renumVN = succ
    , renumVP = succ
    }

addCounts :: OBJ Count -> Renumbering -> Renumbering
addCounts c r = Renumbering
    { renumV  = addc (objLocations c) . renumV r
    , renumVT = addc (objTexCoords c) . renumVT r
    , renumVN = addc (objNormals c) . renumVN r
    , renumVP = renumVP r -- TODO
    }

addc :: Count x -> Int -> Int
addc (Count c) x = c + x

renumTriple :: Renumbering -> RefTriple -> RefTriple
renumTriple r (RefTriple v t n) = RefTriple (renumV r v) (renumVT r <$> t) (renumVN r <$> n)

applyRenumbering :: MonadState Renumbering m => ObjBuilder m -> ObjBuilder m
applyRenumbering builder = builder
    { face = \ts -> do
        r <- get
        face builder $ map (renumTriple r) ts
    , line = \ts -> do
        r <- get
        line builder $ map (renumTriple r) ts
    , surf = \u0 u1 v0 v1 ts -> do
        r <- get
        surf builder u0 u1 v0 v1 $ map (renumTriple r) ts
    }