path: root/src/Wavefront.hs

{-# LANGUAGE NondecreasingIndentation #-}
{-# LANGUAGE OverloadedStrings        #-}
module Wavefront where

import qualified Data.ByteString.Lazy.Char8 as L
import qualified Data.ByteString.Char8 as S
import Data.ByteString.Internal as BS
import Data.Char
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.ByteString.Lex.Fractional as F
import Data.ByteString.Lex.Integral as I

data ObjBuilder m = ObjBuilder
    { vertex           :: [Double] -> m ()
    , vertexT          :: [Double] -> m ()
    , vertexN          :: [Double] -> m ()
    , vertexP          :: [Double] -> m ()
    , face             :: [RefTriple] -> m ()
    , cstype           :: Bool -> CSType -> m ()
    , curv2            :: [Int] -> m ()
    , curv             :: Double -> Double -> [Int] -> m ()
    , parm             :: ParamSpec -> [Double] -> m ()
    , specialPoints    :: [Int] -> m ()
    , endFreeForm      :: m ()
    , ctech            :: CurveSamplingSpec -> m ()
    , stech            :: SurfaceSamplingSpec -> m ()
    , deg              :: [Int] -> m ()
    , surf             :: Double -> Double -> Double -> Double -> [RefTriple] -> m ()
    , trim             :: [CurveSpec] -> m ()
    , hole             :: [CurveSpec] -> m ()
    , specialCurves    :: [CurveSpec] -> m ()
    , equivalentCurves :: [EmbeddedCurve] -> m ()
    , groups           :: [S.ByteString] -> m ()
    , smoothingGroup   :: Int -> m ()
    , mergingGroup     :: Int -> Double -> m ()
    , usemtl           :: S.ByteString -> m ()
    , deprecated_cdc   :: [Int] -> m ()
    , deprecated_bzp   :: [Int] -> m ()
    , mtllib           :: [S.ByteString] -> m ()
    , objectName       :: S.ByteString -> m ()
    , bmat             :: ParamSpec -> [Double] -> m ()
    , badToken         :: L.ByteString -> m ()
    }

nullBuilder :: Applicative m => ObjBuilder m
nullBuilder = ObjBuilder
    { vertex           = \vs -> pure ()
    , vertexT          = \vs -> pure ()
    , vertexN          = \vs -> pure ()
    , vertexP          = \vs -> pure ()
    , face             = \is -> pure ()
    , cstype           = \isRat typ -> pure ()
    , curv2            = \is -> pure ()
    , curv             = \u0 v0 is -> pure ()
    , parm             = \uv is -> pure ()
    , specialPoints    = \is -> pure ()
    , endFreeForm      = pure ()
    , ctech            = \approx -> pure ()
    , stech            = \approx -> pure ()
    , deg              = \is -> pure ()
    , surf             = \u0 u1 v0 v1 ts -> pure ()
    , trim             = \ss -> pure ()
    , hole             = \ss -> pure ()
    , specialCurves    = \ss -> pure ()
    , equivalentCurves = \ccs -> pure ()
    , groups           = \gs -> pure ()
    , smoothingGroup   = \sg -> pure ()
    , mergingGroup     = \mg δ -> pure ()
    , usemtl           = \mtl -> pure ()
    , deprecated_cdc   = \is -> pure ()
    , deprecated_bzp   = \is -> pure ()
    , mtllib           = \fns -> pure ()
    , objectName       = \obn -> pure ()
    , bmat             = \uv fs -> pure ()
    , badToken         = \bs -> pure ()
    }


data CurveSamplingSpec
    -- ctech cparm
    = UniformSubdivision
        { divisionsPerCurveDegree :: Double -- ^ This really ought to be an integer but
                                            -- but examples show floats.  The only way
                                            -- it makes sense as a float is if we are to
                                            -- convert to an integer *after* multiplying
                                            -- by the curve degree.
        }
    -- ctech cspace
    | MaxLengthPolygonal { maxPolygonEdgeLength :: Double }
    -- ctech curv
    | CurvatureBasedPolygon { maxDistanceToCurve :: Double, maximumDegreesPerSample :: Double }
 deriving (Eq,Show)

data SurfaceSamplingSpec
    -- stech cparma ures vres
    = UniformIsoparametric { uDivisionsPerDegree :: Double, vDivisionsPerDegree :: Double }
    -- stech cparmb uvres
    | UniformAfterTrimming { uvDivisionsPerDegree :: Double }
    -- stech cspace maxlength
    | MaxLengthPolytopal { maxPolytopEdgeLength :: Double }
    -- stech curv maxdist maxangle
    | CurvatureBasedPolytope { maxDistanceToSurface :: Double, maxDegreesPerCorner :: Double }
 deriving (Eq,Show)


data ObjState = ObjState
    {
    }

newtype ObjConfig = ObjConfig
    { cfgSubst :: IntMap L.ByteString
    }

-- consChunk :: S.ByteString -> L.ByteString -> L.ByteString
-- consChunk c bs = L.fromChunks (c : L.toChunks bs)

reappend :: ByteString -> ByteString -> Maybe ByteString
reappend a b =
    let (ap,ao,al) = BS.toForeignPtr a
        (bp,bo,bl) = BS.toForeignPtr b
    in if ap == bp && ao+al == bo
        then Just $ BS.PS ap ao (al+bl)
        else Nothing

reconsChunk :: S.ByteString -> L.ByteString -> L.ByteString
reconsChunk b bs = case L.toChunks bs of
    (c:cs) -> case reappend b c of
        Just x  -> L.fromChunks (x:cs)
        Nothing -> L.fromChunks (b:c:cs)
    _      -> L.fromChunks [b]

findToken :: ObjConfig -> L.ByteString -> L.ByteString
findToken (ObjConfig args) bs = case L.dropWhile (\c -> isSpace c || c=='\\') bs of
  cs -> case L.uncons cs of
    Just ('#',comment) -> findToken (ObjConfig args) $ L.drop 1 $ L.dropWhile (/='\n') comment
    Just ('$',ref)     -> case L.splitAt 5 ref of
                            (refp,ds) -> case I.readDecimal (L.toStrict refp) of
                                Just (i,es) -> case IntMap.lookup i args of
                                    Just val -> val <> reconsChunk es ds
                                    _        -> reconsChunk es ds
                                _ -> cs
    Just _             -> cs
    Nothing            -> L.empty

{-

     1	x   bevel
     2	    bmat
     3	    bzp
     4	x   call
     5	    cdc
     6	x   c_interp
     7	    con
     8	x   csh    -- for all except these,
     9	    cstype -- Two chars suffice to distinguish
    10	    ctech
    11	    curv2 -- for all except these,
    12	    curv  -- Two chars suffice to distinguish
    13	    deg
    14	x   d_interp
    15	    end
    16	    f
    17	    g
    18	    hole
    19	x   lod
    20	x   maplib
    21	    mg
    22	    mtllib
    23	    o
    24	    parm
    25	    s
    26	    scrv
    27	x   shadow_obj
    28	    sp
    29	    stech -- for all except these,
    30	x   step  -- Two chars suffice to distinguish
    31	    surf
    32	x   trace_obj -- for all except these,
    33	    trim      -- Two chars suffice to distinguish
    34	x   usemap -- for all except these,
    35	    usemtl -- Two chars suffice to distinguish
    36	    v
    37	    vn
    38	    vp
    39	    vt

-}

nextToken :: (L.ByteString -> L.ByteString) -> L.ByteString -> L.ByteString
nextToken tok bs = tok $ L.dropWhile (not . isSpace) bs

parseFloats tok bs cont = case L.splitAt 10 (tok bs) of
    (ds,bs') -> case F.readSigned F.readExponential (L.toStrict ds) of
        Just (x,b) -> parseFloats tok (reconsChunk b bs') (cont . (x :))
        Nothing    -> cont [] bs

parseFloatsN 0 _   bs cont = cont [] bs
parseFloatsN n tok bs cont = case L.splitAt 10 (tok bs) of
    (ds,bs') -> case F.readSigned F.readExponential (L.toStrict ds) of
        Just (x,b) -> parseFloatsN (n-1) tok (reconsChunk b bs') (cont . (x :))
        Nothing    -> cont [] bs

parseInts tok bs cont = case L.splitAt 5 (tok bs) of
    (ds,bs') -> case I.readSigned I.readDecimal (L.toStrict ds) of
        Just (x,b) -> parseInts tok (reconsChunk b bs') (cont . (x :))
        Nothing    -> cont [] bs

parseIntsN 0 tok bs cont = cont [] bs
parseIntsN n tok bs cont = case L.splitAt 5 (tok bs) of
    (ds,bs') -> case I.readSigned I.readDecimal (L.toStrict ds) of
        Just (x,b) -> parseIntsN (n-1) tok (reconsChunk b bs') (cont . (x :))
        Nothing    -> cont [] bs

-- Optimize me
parseTriples :: (L.ByteString -> L.ByteString) -> L.ByteString -> ([RefTriple] -> L.ByteString -> b) -> b
parseTriples tok bs cont = case L.splitAt 17 (tok bs) of
    (ds,bs') -> case I.readSigned I.readDecimal (L.toStrict ds) of

        Just (v,b) -> case S.splitAt 1 b of

            ("/",ds') -> case I.readSigned I.readDecimal ds' of

                Just (vt,c) -> case S.splitAt 1 c of
                    ("/",ds'') -> case I.readSigned I.readDecimal ds'' of
                        Just (vn,d) -> parseTriples tok (reconsChunk d bs') $ cont . (RefTriple v (Just vt) (Just vn) :)
                        Nothing     -> parseTriples tok (reconsChunk ds'' bs') $ cont . (RefTriple v (Just vt) Nothing :)

                    _          -> parseTriples tok (reconsChunk c bs') $ cont . (RefTriple v (Just vt) Nothing :)

                Nothing -> case S.splitAt 1 ds' of
                    ("/",ds'') -> case I.readSigned I.readDecimal ds'' of
                        Just (vn,d) -> parseTriples tok (reconsChunk d bs') $ cont . (RefTriple v Nothing (Just vn) :)
                        Nothing     -> parseTriples tok (reconsChunk ds'' bs') $ cont . (RefTriple v Nothing Nothing :)

                    _          -> parseTriples tok (reconsChunk ds' bs') $ cont . (RefTriple v Nothing Nothing :)


            _         -> parseTriples tok (reconsChunk b bs') $ cont . (RefTriple v Nothing Nothing :)

        Nothing -> cont [] bs

parseCurveSpecs :: (L.ByteString -> L.ByteString) -> L.ByteString -> ([CurveSpec] -> L.ByteString -> b) -> b
parseCurveSpecs tok bs cont = parseFloatsN 2 tok bs $ \fs bs' -> case fs of
    (u0:u1:_) -> do
        parseIntsN 1 tok bs' $ \is bs'' -> case is of
            (i:_) -> parseCurveSpecs tok bs'' $ cont . (CurveSpec u0 u1 i :)
            _     -> cont [] bs''

    _ -> cont [] bs'


parseCurveSpecsN :: Int -> (L.ByteString -> L.ByteString) -> L.ByteString -> ([CurveSpec] -> L.ByteString -> b) -> b
parseCurveSpecsN 0 tok bs cont = cont [] bs
parseCurveSpecsN n tok bs cont = parseFloatsN 2 tok bs $ \fs bs' -> case fs of
    (u0:u1:_) -> do
        parseIntsN 1 tok bs' $ \is bs'' -> case is of
            (i:_) -> parseCurveSpecsN (n-1) tok bs'' $ cont . (CurveSpec u0 u1 i :)
            _     -> cont [] bs''

    _ -> cont [] bs'

parseEmbeddedCurves :: (L.ByteString -> L.ByteString) -> L.ByteString -> ([EmbeddedCurve] -> L.ByteString -> b) -> b
parseEmbeddedCurves tok bs cont = parseIntsN 1 tok bs $ \is bs' -> case is of
    (sref:_) -> do
        parseCurveSpecsN 1 tok bs' $ \cs bs'' -> case cs of
            (c:_) -> parseEmbeddedCurves tok bs'' $ cont . (EmbeddedCurve sref c :)
            _     -> cont [] bs''

    _ -> cont [] bs'

data CSType = Bmatrix | Bezier | Bspline | Cardinal | Taylor
 deriving (Eq,Ord,Show,Enum)

data ParamSpec = ParamU | ParamV
 deriving (Eq,Ord,Show,Enum)

data RefTriple = RefTriple
        { refV :: Int
        , refT :: Maybe Int
        , refN :: Maybe Int
        }
-- data RefTriple = RefTriple Int (Maybe Int) (Maybe Int)
 deriving (Eq,Ord,Show)

data CurveSpec = CurveSpec
    { curveStart :: Double
    , curveEnd   :: Double
    , curveRef   :: Int
    }
 deriving (Eq,Ord,Show)

data EmbeddedCurve = EmbeddedCurve
    { curveSurfaceRef :: Int
    , embeddedCurve :: CurveSpec
    }
 deriving (Eq,Ord,Show)

lengthLessThan :: Int -> L.ByteString -> Bool
lengthLessThan n bs =
    foldr (\c ret ac -> let m = S.length c in if ac <= m then False else ret $! ac - m)
          (const True)
          (L.toChunks bs)
          n

parseOBJ :: Monad m => ObjBuilder m -> ObjConfig -> L.ByteString -> m ()
parseOBJ builder args bs0
    | lengthLessThan 2 bs = return ()
    | isSpace (L.index bs 1)  = case L.head bs of
        'f' -> parseT face 2
        'g' -> case L.break (=='\n') $ L.drop 1 bs of -- Newline required to terminate group name list.
                (gn,bs') -> do
                    groups builder (map L.toStrict $ L.words gn)
                    parseOBJ builder args bs'
        's' -> case next 1 bs of
                tok -> parseOffOrNumber tok $ \sg bs' -> do
                            smoothingGroup builder sg
                            parseOBJ builder args bs'
        'v' -> parseV vertex 2
        'o' -> -- o object-name
               case L.break (=='\n') $ findToken args $ L.drop 1 bs of
                (objn,bs') -> do
                    objectName builder (L.toStrict objn)
                    parseOBJ builder args bs'
        _   -> badToken builder bs
    | otherwise = case L.take 2 bs of
        "vt" -> parseV vertexT 3
        "vn" -> parseV vertexN 3
        "vp" -> parseV vertexP 3
        "bz" -> parseI deprecated_bzp 4 -- bzp
        "bm" -> parseUV (next 2 bs) $ \uv bs' -> do -- bmat
                    parseFloats (findToken args) bs' $ \vs bs'' -> do
                        bmat builder uv vs
                        parseOBJ builder args bs''

        "cd" -> parseI deprecated_cdc 4 -- cdc
        "co" -> -- con
                parseEmbeddedCurves (findToken args) (next 2 bs) $ \ss bs' -> do
                    equivalentCurves builder ss
                    parseOBJ builder args bs'
        "cs" -> -- cstype
                let parseRat = parseChar 'r'
                    parseTyp tok cont | lengthLessThan 3 tok = badToken builder tok
                                      | otherwise = case L.index tok 2 of
                        'a' -> cont Bmatrix  $ next 3 tok
                        'z' -> cont Bezier   $ next 3 tok
                        'p' -> cont Bspline  $ next 3 tok
                        'r' -> cont Cardinal $ next 3 tok
                        'y' -> cont Taylor   $ next 3 tok
                        _   -> badToken builder tok
                in parseRat (next 2 bs) $ \isRat bs' -> do
                    parseTyp bs' $ \typ bs'' -> do
                        cstype builder isRat typ
                        parseOBJ builder args bs''
        "ct" -> -- ctech
                let tok = next 2 bs
                in if lengthLessThan 2 tok
                    then badToken builder tok
                    else case L.index tok 1 of
                        'p' -> -- cparm
                               parseFloats (findToken args) (next 2 tok) $ \is bs' -> do
                                    let x:_ = is ++ [0]
                                    ctech builder (UniformSubdivision x)
                                    parseOBJ builder args bs'
                        's' -> -- cspace
                               parseFloats (findToken args) (next 2 tok) $ \fs bs' -> do
                                    let x:_ = fs ++ [1.0]
                                    ctech builder (MaxLengthPolygonal x)
                                    parseOBJ builder args bs'
                        'u' -> -- curv
                               parseFloats (findToken args) (next 2 tok) $ \fs bs' -> do
                                    let δ:θ:_ = fs ++ repeat 1.0
                                    ctech builder (CurvatureBasedPolygon δ θ)
                                    parseOBJ builder args bs'
                        _   -> badToken builder tok
        "cu" -> if lengthLessThan 5 bs
                    then badToken builder bs
                    else if L.index bs 4 == '2'
                            then parseI curv2 5 -- curv2
                            else do -- curv
                                parseFloatsN 2 (findToken args) (L.drop 4 bs) $ \vs bs' ->
                                    parseInts (findToken args) bs' $ \is bs'' -> do
                                        let u0:v0:_ = vs ++ repeat 0.0
                                        curv builder u0 v0 is
                                        parseOBJ builder args bs''
        "de" -> parseI deg 3
        "en" -> do endFreeForm builder
                   parseOBJ builder args (next 2 bs)
        "ho" -> -- hole
                parseCurveSpecs (findToken args) (next 2 bs) $ \ss bs' -> do
                    hole builder ss
                    parseOBJ builder args bs'
        "mg" -> case next 2 bs of
                tok -> parseOffOrNumber tok $ \mg bs' -> do
                            parseFloatsN 1 (findToken args) bs' $ \fs bs'' -> do
                                mergingGroup builder mg (head $ fs ++ [0])
                                parseOBJ builder args bs''
        "pa" -> parseUV (next 2 bs) $ \uv bs' -> do
                    parseFloats (findToken args) bs' $ \vs bs'' -> do
                       parm builder uv vs
                       parseOBJ builder args bs''
        "sc" -> -- scrv
                parseCurveSpecs (findToken args) (next 2 bs) $ \ss bs' -> do
                    specialCurves builder ss
                    parseOBJ builder args bs'
        "sp" -> parseI specialPoints 3
        "st" -> -- stech
                let tok = next 2 bs
                in if lengthLessThan 2 tok
                    then badToken builder tok
                    else case L.index tok 1 of
                        'p' -> -- cparma/cparmb
                               if lengthLessThan 6 tok
                                then badToken builder tok
                                else if L.index tok 5 == 'b'
                                        then -- cparmb
                                           parseFloats (findToken args) (next 5 tok) $ \is bs' -> do
                                                let x:_ = is ++ [0]
                                                stech builder (UniformAfterTrimming x)
                                                parseOBJ builder args bs'
                                        else -- cparma
                                           parseFloats (findToken args) (next 5 tok) $ \is bs' -> do
                                                let x:y:_ = is ++ [0]
                                                stech builder (UniformIsoparametric x y)
                                                parseOBJ builder args bs'
                        's' -> -- cspace
                               parseFloats (findToken args) (next 2 tok) $ \fs bs' -> do
                                    let x:_ = fs ++ [1.0]
                                    stech builder (MaxLengthPolytopal x)
                                    parseOBJ builder args bs'
                        'u' -> -- curv
                               parseFloats (findToken args) (next 2 tok) $ \fs bs' -> do
                                    let δ:θ:_ = fs ++ repeat 1.0
                                    stech builder (CurvatureBasedPolytope δ θ)
                                    parseOBJ builder args bs'
                        _   -> badToken builder tok
        "su" -> -- surf
                parseFloatsN 4 (findToken args) (next 2 bs) $ \fs bs' -> do
                    parseTriples (findToken args) bs' $ \ts bs'' -> do
                        let u0:u1:v0:v1:_ = fs ++ repeat 0
                        surf builder u0 u1 v0 v1 ts
                        parseOBJ builder args bs''
        "tr" -> -- trim
                parseCurveSpecs (findToken args) (next 2 bs) $ \ss bs' -> do
                    trim builder ss
                    parseOBJ builder args bs'
        "us" -> -- usemtl
                case L.break isSpace $ next 2 bs of
                    (mtl,bs') -> do
                        usemtl builder (L.toStrict mtl)
                        parseOBJ builder args bs'
        "mt" -> -- mtllib
                case L.break (=='\n') $ next 2 bs of
                    (fnn,bs') -> do
                        let slurp fnn = case L.break (=='.') fnn of
                                (fn,ext) | L.null (L.drop 1 ext) -> if L.null fn then []
                                                                                 else [fn]
                                         | ".mtl" <- L.take 4 ext ->
                                                if L.all isSpace (L.take 1 $ L.drop 4 ext)
                                                    then (fn <> ".mtl") : slurp (findToken args $ L.drop 4 ext)
                                                    else let f:fs = slurp (L.drop 3 ext) in (fn <> L.take 3 ext <> f) : fs
                                         | otherwise -> let (f:fs) = slurp (L.drop 1 ext)
                                                        in (fn <> L.take 1 ext <> f) : fs
                        mtllib builder (map L.toStrict $ slurp fnn)
                        parseOBJ builder args bs'
        -- TODO: call,csh,step,c_interp,d_interp,lod,shadow_obj,trace_obj,bevel
        _    -> badToken builder bs
    where
        bs = findToken args bs0
        next n xs = nextToken (findToken args) $ L.drop n xs
        parseChar c tok cont = case L.uncons tok of
                        Just (x,cs) | x==c -> cont True $ next 0 cs
                        _                  -> cont False tok
        parseUV tok cont = parseChar 'u' tok $ \isU bs' -> do
            cont (if isU then ParamU else ParamV)
                 (if isU then bs' else L.drop 1 bs')
        parseV build n = do
                parseFloats (findToken args) (L.drop n bs) $ \vs bs' -> do
                build builder vs
                parseOBJ builder args bs'
        parseI build n = do
                parseInts (findToken args) (L.drop n bs) $ \vs bs' -> do
                build builder vs
                parseOBJ builder args bs'
        parseT build n = do
                parseTriples (findToken args) (L.drop n bs) $ \vs bs' -> do
                build builder vs
                parseOBJ builder args bs'
        parseOffOrNumber tok cont = parseIntsN 1 (findToken args) tok $ \is bs' -> do
            let (sg,bs'') = case is of
                        (i:_)                    -> (i,bs')
                        _ | lengthLessThan 2 tok -> (0,tok)
                        _                        -> (if L.index tok 1 == 'f' then 0 else 1, next 1 tok)
            cont sg bs''