Support zipCount primitive to access gl_VertexID.

2 files changed, 37 insertions, 19 deletions

diff --git a/lc/Builtins.lc b/lc/Builtins.lc
index c2e2e7a9..fe70f277 100644
--- a/lc/Builtins.lc
+++ b/lc/Builtins.lc
@@ -583,4 +583,6 @@ textureBuffer :: TextureBuffer -> Int -> Float
 -- todo: remove
 accumulationContext x = x
 
-
+-- gl_VertexID
+count :: Int -> PrimitiveStream a (HList '[Int]) -- TODO: implement this case
+zipCount :: PrimitiveStream pr (HList a) -> PrimitiveStream pr (HList (Int:a))
diff --git a/src/LambdaCube/Compiler/CoreToIR.hs b/src/LambdaCube/Compiler/CoreToIR.hs
index 2ad8d26d..7f4f9561 100644
--- a/src/LambdaCube/Compiler/CoreToIR.hs
+++ b/src/LambdaCube/Compiler/CoreToIR.hs
@@ -134,13 +134,13 @@ getCommands backend e = case e of
 
     A3 "Accumulate" actx (getFragmentShader -> (frag, getFragFilter -> (ffilter, x1))) fbuf -> case x1 of
 
-      A3 "foldr" (A0 "++") (A0 "Nil") (A2 "map" (EtaPrim3 "rasterizePrimitive" ints rctx) (getVertexShader -> (vert, input_))) -> mdo
+      A3 "foldr" (A0 "++") (A0 "Nil") (A2 "map" (EtaPrim3 "rasterizePrimitive" ints rctx) ex@(getVertexShader -> (vert, input_, hasCount))) -> mdo
 
         let
             (vertexInput, pUniforms, vertSrc, fragSrc) = case backend of
               -- disabled DX11 codegen, due to it's incomplete
               --IR.DirectX11 -> genHLSLs backend (compRC' rctx) ints vert frag ffilter
-              _ -> genGLSLs backend (compRC' rctx) ints vert frag ffilter
+              _ -> genGLSLs backend (compRC' rctx) ints hasCount vert frag ffilter
 
             pUniforms' = snd <$> Map.filter ((\case UTexture2D{} -> False; _ -> True) . fst) pUniforms
 
@@ -176,8 +176,7 @@ getCommands backend e = case e of
 
         (smpBindings, txtCmds) <- mconcat <$> traverse (uncurry getRenderTextureCommands) (Map.toList $ fst <$> pUniforms)
 
-        (renderCommand,input) <- case input_ of
-            A2 "fetch" (EString slotName) attrs -> do
+        let grokFetch slotName attrs = do
                 i <- IR.RenderSlot <$> addL' slotLens slotName (flip mergeSlot) IR.Slot
                     { IR.slotName       = slotName
                     , IR.slotUniforms   = IR.programUniforms prg
@@ -193,6 +192,12 @@ getCommands backend e = case e of
                   , IR.slotStreams  = IR.slotStreams a <> IR.slotStreams b
                   , IR.slotPrograms = IR.slotPrograms a <> IR.slotPrograms b
                   }
+
+        (renderCommand,input) <- case input_ of
+            A2 "fetch" (EString slotName) attrs -> grokFetch slotName attrs
+            A1 "zipCount" (A2 "fetch" (EString slotName) attrs) -> do
+                (rc,inp) <- grokFetch slotName attrs
+                return (rc,inp)
             A1 "fetchArrays" (unzip . compAttributeValue -> (tys, values)) -> do
                 i <- IR.RenderStream <$> addL streamLens IR.StreamData
                     { IR.streamData       = Map.fromList $ zip names values
@@ -302,11 +307,15 @@ getFragFilter :: ExpTV -> (Maybe ExpTV, ExpTV)
 getFragFilter (A2 "map" (EtaPrim2 "filterFragment" p) x) = (Just p, x)
 getFragFilter x = (Nothing, x)
 
-getVertexShader :: ExpTV -> ((Maybe ExpTV, Ty), ExpTV)
-getVertexShader (A2 "map" (EtaPrim2 "mapPrimitive" f@(etaReds -> Just (_, o))) x) = ((Just f, tyOf o), x)
+getVertexShader :: ExpTV -> ((Maybe ExpTV, Ty), ExpTV, Bool)
+getVertexShader xx@(A2 "map" (EtaPrim2 "mapPrimitive" f@(etaReds -> Just (_, o))) x) =
+   let hasCount = case x of
+        A1 "zipCount" _ -> True
+        _               -> False
+   in ((Just f, tyOf o), x, hasCount)
 --getVertexShader (A2 "map" (EtaPrim2 "mapPrimitive" f) x) = error $ "gff: " ++ show (case f of ExpTV x _ _ -> x) --ppShow (mapVal unFunc' f)
 --getVertexShader x = error $ "gf: " ++ ppShow x
-getVertexShader x = ((Nothing, getPrim' $ tyOf x), x)
+getVertexShader x = ((Nothing, getPrim' $ tyOf x), x, False {- hasCount: TODO: is this right? -})
 
 getFragmentShader :: ExpTV -> ((Maybe ExpTV, Ty), ExpTV)
 getFragmentShader (A2 "map" (EtaPrim2 "mapFragment" f@(etaReds -> Just (_, o))) x) = ((Just f, tyOf o), x)
@@ -662,18 +671,20 @@ genGLSLs
   :: Backend
   -> Maybe ExpTV
   -> ExpTV
+  -> Bool
   -> (Maybe ExpTV, ExpTV)
   -> (Maybe ExpTV, ExpTV)
   -> Maybe ExpTV
   -> ([[Char]], Uniforms, Doc, Doc)
 genGLSLs backend
-    rp                  -- program point size
-    (ETuple ints)       -- interpolations
-    (vert, tvert)       -- vertex shader
-    (frag, tfrag)       -- fragment shader
-    ffilter             -- fragment filter
+    rp            -- program point size
+    (ETuple ints) -- interpolations
+    hasCount      -- True if vertex shader uses gl_VertexID
+    (vert, tvert) -- vertex shader
+    (frag, tfrag) -- fragment shader
+    ffilter       -- fragment filter
     = ( -- vertex input
-        vertInNames
+        (if hasCount then drop 1 else id) vertInNames
 
       , -- uniforms
         vertUniforms <> fragUniforms
@@ -681,7 +692,7 @@ genGLSLs backend
       , -- vertex shader code
         shader $
            uniformDecls vertUniforms
-        <> [shaderDecl (caseWO "attribute" "in") (text t) (text n) | (n, t) <- zip vertInNames vertIns]
+        <> (if hasCount then drop 1 else id) [shaderDecl (caseWO "attribute" "in") (text t) (text n) | (n, t) <- zip vertInNames vertIns]
         <> vertOutDecls "out"
         <> vertFuncs
         <> [mainFunc $
@@ -711,9 +722,11 @@ genGLSLs backend
         1 -> [caseWO "gl_FragColor" "f0"]
 
     (vertIns, verts) = case vert of
-        Just (etaReds -> Just (xs, ETuple ys)) -> (toGLSLType "3" <$> xs, ys)
+        Just (etaReds -> Just (xs, expr)) -> case expr of
+            ETuple ys -> (toGLSLType "3" <$> xs, ys)
+            _         -> error ("mapPrimitives function body is not a tuple: " ++ ppShow expr)
         Nothing -> ([toGLSLType "4" tvert], [mkTVar 0 tvert])
-        Just wut -> error $ show wut
+        Just wut -> error $ ppShow (etaReds wut, wut)
 
     (fragOuts, frags) = case frag of
         Just (etaReds -> Just (xs, ETuple ys)) -> (toGLSLType "31" . tyOf <$> ys, ys)
@@ -754,7 +767,10 @@ genGLSLs backend
 
     funArgs      = map (("z" ++) . show) [0..]
     shaderNames  = map (("s" ++) . show)  [0..]
-    vertInNames  = map (("vi" ++) . show) [1..length vertIns]
+    vertInNames  = let vns = map (("vi" ++) . show) [1..]
+                       ns | hasCount  = "gl_VertexID" : vns
+                          | otherwise = vns
+                    in zipWith const ns vertIns
     vertOutNames = map (("vo" ++) . show) [1..length vertOuts]
     vertOutNamesWithPosition = "gl_Position": vertOutNames
 
@@ -1537,7 +1553,7 @@ genHLSLs backend
 
     funArgs      = map (("z" ++) . show) [0..]
     shaderNames  = map (("s" ++) . show)  [0..]
-    vertInNames  = map (("vi" ++) . show) [1..length vertIns]
+    vertInNames  = map (("vi" ++) . show) [1..length vertIns] -- TODO: hasCount support?
     vertOutNames = map (("vo" ++) . show) [1..length vertOuts]
     vertOutNamesWithPosition = "gl_Position": vertOutNames