1 files changed, 137 insertions, 0 deletions

diff --git a/src/Wavefront.hs b/src/Wavefront.hs
new file mode 100644
index 0000000..6195eee
--- /dev/null
+++ b/src/Wavefront.hs
@@ -0,0 +1,137 @@
+{-# LANGUAGE FlexibleContexts #-}
+module Wavefront where
+
+import Data.Wavefront
+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 Data.Text.Encoding                   (decodeUtf8)
+import qualified Data.Vector                as Vector
+         ;import Data.Vector                (Vector)
+import qualified Rank2
+
+newtype Count x = Count Int
+
+updateCount :: MonadState s m => (s -> Count x) -> (Count x -> s -> s) -> m ()
+updateCount 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 :: WavefrontOBJ 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 (WavefrontOBJ Count) m)
+objBookKeeping builder = (lift Rank2.<$> builder)
+    { vertex = \xs -> do lift $ vertex builder xs
+                         updateCount objLocations $ \x o -> o { objLocations = x }
+    , vertexT = \xs -> do lift $ vertexT builder xs
+                          updateCount objTexCoords $ \x o -> o { objTexCoords = x }
+    , vertexN = \xs -> do lift $ vertexN builder xs
+                          updateCount objNormals $ \x o -> o { objNormals = x }
+    , points = \xs -> do
+        n <- gets objLocations
+        lift $ points builder $ fixupRef n <$> xs
+    , line = \ts -> do
+        o <- get
+        lift $ line builder $ fixupTriple o <$> ts
+    , face = \ts -> do
+        o <- get
+        lift $ face builder $ fixupTriple o <$> ts
+    }
+
+mkv :: [Double] -> Location
+mkv cs = Location x y z w where (x:y:z:w:_) = map realToFrac cs
+
+mkt :: [Double] -> TexCoord
+mkt cs = TexCoord x y z where (x:y:z:_) = map realToFrac cs
+
+mkn :: [Double] -> Normal
+mkn cs = Normal x y z where (x:y:z:_) = map realToFrac cs
+
+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 (c, d) m => (c -> c) -> m ()
+modifyFirst = modify' . first
+
+buildOBJ :: ObjBuilder (State (WavefrontOBJ DList,Element ()))
+buildOBJ = (nullBuilder $ pure ())
+    { 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 . fst)
+        modifyFirst $ \o -> o { objMtlLibs = libs `DList.append` DList.fromList l }
+    , groups = \xs -> do
+        let g = map decodeUtf8 xs
+        modify' $ second $ \e -> e { elGroups = g }
+    , objectName = \x -> do
+        let o = decodeUtf8 x
+        modify' $ second $ \e -> e { elObject = Just o }
+    , usemtl = \x -> do
+        let mtl = decodeUtf8 x
+        modify' $ second $ \e -> e { elMtl = Just mtl }
+    , smoothingGroup = \x -> when (x > 0) $ do
+        modify $ second $ \e -> e { elSmoothingGroup = fromIntegral x }
+    }
+
+parse :: L.ByteString -> WavefrontOBJ Vector
+parse bs = Rank2.fmap (Vector.fromList . DList.toList) obj
+ where
+    b = objBookKeeping buildOBJ
+    ls = WavefrontOBJ
+            { objLocations = DList.empty
+            , objTexCoords = DList.empty
+            , objNormals   = DList.empty
+            , objPoints    = DList.empty
+            , objLines     = DList.empty
+            , objFaces     = DList.empty
+            , objMtlLibs   = DList.empty
+            }
+    c = Rank2.fmap (const $ Count 0) ls :: WavefrontOBJ Count
+    el = Element
+            { elObject         = Nothing
+            , elGroups         = []
+            , elMtl            = Nothing
+            , elSmoothingGroup = 0
+            , elValue          = ()
+            }
+    substvars = ObjConfig IntMap.empty
+    (obj,_) = execState (runStateT (parseOBJ b substvars bs) c) (ls,el)