Compile: replicate1i

1 files changed, 25 insertions, 3 deletions

diff --git a/src/Compile.hs b/src/Compile.hs
index 3b02c60..7a846d4 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -20,7 +20,7 @@ import Data.Set (Set)
 import Data.Some
 import qualified Data.Vector as V
 import Foreign
-import System.IO (hPutStrLn, stderr)
+-- import System.IO (hPutStrLn, stderr)
 
 import Prelude hiding ((^))
 import qualified Prelude
@@ -664,12 +664,28 @@ compile' env = \case
     return (CELit name)
 
   EReplicate1Inner _ elen earg -> do
+    let STArr n t = typeOf earg
     lenname <- genName' "replen"
     emit . SVarDecl True (repSTy tIx) lenname =<< compile' env elen
     argname <- genName' "reparg"
     emit . SVarDecl True (repSTy (typeOf earg)) argname =<< compile' env earg
 
-    _
+    shszname <- genName' "shsz"
+    emit $ SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n argname)
+
+    resname <- allocArray "rep" (SS n) t
+                 (CEBinop (CELit shszname) "*" (CELit lenname))
+                 ([CELit (argname ++ ".buf->sh[" ++ show i ++ "]") | i <- [0 .. fromSNat n - 1]]
+                  ++ [CELit lenname])
+
+    ivar <- genName' "i"
+    jvar <- genName' "j"
+    emit $ SLoop (repSTy tIx) ivar (CELit "0") (CELit shszname) $
+             pure $ SLoop (repSTy tIx) jvar (CELit "0") (CELit lenname) $
+               pure $ SAsg (resname ++ ".buf->a[" ++ ivar ++ " * " ++ lenname ++ " + " ++ jvar ++ "]")
+                           (CELit (argname ++ ".buf->a[" ++ ivar ++ "]"))
+
+    return (CELit resname)
 
   -- EMaximum1Inner _ e -> error "TODO" -- EMaximum1Inner ext (compile' e)
 
@@ -737,7 +753,13 @@ compile' env = \case
   EPlus{} -> error "Compile: monoid operations should have been eliminated"
   EOneHot{} -> error "Compile: monoid operations should have been eliminated"
 
-  _ -> error "Compile: not implemented"
+  EFold1Inner{} -> error "Compile: not implemented: EFold1Inner"
+  EMaximum1Inner{} -> error "Compile: not implemented: EMaximum1Inner"
+  EMinimum1Inner{} -> error "Compile: not implemented: EMinimum1Inner"
+  EIdx1{} -> error "Compile: not implemented: EIdx1"
+  ECustom{} -> error "Compile: not implemented: ECustom"
+  EWith{} -> error "Compile: not implemented: EWith"
+  EAccum{} -> error "Compile: not implemented: EAccum"
 
 data Increment = Increment | Decrement
   deriving (Show)