Compile: compileAssign helper function

1 files changed, 27 insertions, 45 deletions

diff --git a/src/Compile.hs b/src/Compile.hs
index d9cfd95..7e6fd5c 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -278,6 +278,7 @@ genId :: CompM Int
 genId = state $ \s -> (csNextId s, s { csNextId = csNextId s + 1 })
 
 genName' :: String -> CompM String
+genName' "" = genName
 genName' prefix = (prefix ++) . show <$> genId
 
 genName :: CompM String
@@ -490,9 +491,7 @@ compile' env = \case
     return $ CELit var
 
   ELet _ rhs body -> do
-    e <- compile' env rhs
-    var <- genName
-    emit $ SVarDecl True (repSTy (typeOf rhs)) var e
+    var <- compileAssign "" env rhs
     rete <- compile' (Const var `SCons` env) body
     incrementVarAlways Decrement (typeOf rhs) var
     return rete
@@ -608,8 +607,7 @@ compile' env = \case
     return (CEStruct strname [("buf", CEAddrOf (CELit tldname))])
 
   EBuild _ n esh efun -> do
-    shname <- genName' "sh"
-    emit . SVarDecl True (repSTy (typeOf esh)) shname =<< compile' env esh
+    shname <- compileAssign "sh" env esh
     shsizename <- genName' "shsz"
     emit $ SVarDecl True "size_t" shsizename (compileShapeSize n shname)
 
@@ -636,9 +634,7 @@ compile' env = \case
 
   ESum1Inner _ e -> do
     let STArr (SS n) t = typeOf e
-    e' <- compile' env e
-    argname <- genName' "sumarg"
-    emit $ SVarDecl True (repSTy (STArr (SS n) t)) argname e'
+    argname <- compileAssign "sumarg" env e
 
     shszname <- genName' "shsz"
     -- This n is one less than the shape of the thing we're querying, which is
@@ -677,10 +673,8 @@ compile' env = \case
 
   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
+    lenname <- compileAssign "replen" env elen
+    argname <- compileAssign "reparg" env earg
 
     shszname <- genName' "shsz"
     emit $ SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n argname)
@@ -707,9 +701,7 @@ compile' env = \case
 
   EIdx0 _ e -> do
     let STArr _ t = typeOf e
-    e' <- compile' env e
-    arrname <- genName
-    emit $ SVarDecl True (repSTy (STArr SZ t)) arrname e'
+    arrname <- compileAssign "" env e
     name <- genName
     emit $ SVarDecl True (repSTy t) name
               (CEIndex (CEPtrProj (CEProj (CELit arrname) "buf") "xs") (CELit "0"))
@@ -720,10 +712,10 @@ compile' env = \case
 
   EIdx _ earr eidx -> do
     let STArr n t = typeOf earr
-    arrname <- genName' "ixarr"
-    idxname <- genName' "ixix"
-    emit . SVarDecl True (repSTy (typeOf earr)) arrname =<< compile' env earr
-    when (fromSNat n > 0) $ emit . SVarDecl True (repSTy (typeOf eidx)) idxname =<< compile' env eidx
+    arrname <- compileAssign "ixarr" env earr
+    idxname <- if fromSNat n > 0  -- prevent an unused-varable warning
+                 then compileAssign "ixix" env eidx
+                 else return ""  -- won't be used in this case
     resname <- genName' "ixres"
     emit $ SVarDecl True (repSTy t) resname (CEIndex (CELit (arrname ++ ".buf->xs")) (toLinearIdx n arrname idxname))
     incrementVarAlways Decrement (STArr n t) arrname
@@ -733,8 +725,7 @@ compile' env = \case
     let STArr n _ = typeOf e
         t = tTup (sreplicate n tIx)
     _ <- emitStruct t
-    name <- genName
-    emit . SVarDecl True (repSTy (typeOf e)) name =<< compile' env e
+    name <- compileAssign "" env e
     resname <- genName
     emit $ SVarDecl True (repSTy t) resname (compileShapeQuery n name)
     incrementVarAlways Decrement (typeOf e) name
@@ -749,19 +740,13 @@ compile' env = \case
     e' <- compile' env e
     compileOpGeneral op e'
 
-  ECustom _ t1 t2 _ earg _ _ e1 e2 -> do
-    e1' <- compile' env e1
-    name1 <- genName
-    emit $ SVarDecl True (repSTy t1) name1 e1'
-    e2' <- compile' env e2
-    name2 <- genName
-    emit $ SVarDecl True (repSTy t2) name2 e2'
+  ECustom _ _ _ _ earg _ _ e1 e2 -> do
+    name1 <- compileAssign "" env e1
+    name2 <- compileAssign "" env e2
     compile' (Const name2 `SCons` Const name1 `SCons` SNil) earg
 
   EWith _ t e1 e2 -> do
-    e1' <- compile' env e1
-    name1 <- genName
-    emit $ SVarDecl True (repSTy (typeOf e1)) name1 e1'
+    name1 <- compileAssign "" env e1
 
     mcopy <- copyForWriting t name1
     accname <- genName' "accum"
@@ -772,17 +757,9 @@ compile' env = \case
     return $ CEStruct (repSTy (STPair (typeOf e2) t)) [("a", e2'), ("b", CELit accname)]
 
   EAccum _ t prj eidx eval eacc -> do
-    eidx' <- compile' env eidx
-    nameidx <- genName
-    emit $ SVarDecl True (repSTy (typeOf eidx)) nameidx eidx'
-
-    eval' <- compile' env eval
-    nameval <- genName
-    emit $ SVarDecl True (repSTy (typeOf eval)) nameval eval'
-
-    eacc' <- compile' env eacc
-    nameacc <- genName
-    emit $ SVarDecl True (repSTy (typeOf eacc)) nameacc eacc'
+    nameidx <- compileAssign "acidx" env eidx
+    nameval <- compileAssign "acval" env eval
+    nameacc <- compileAssign "acac" env eacc
 
     let accumRef :: STy a -> SAcPrj p a b -> String -> String -> String
         accumRef _ SAPHere v _ = v
@@ -849,6 +826,13 @@ compile' env = \case
   EFold1Inner{} -> error "Compile: not implemented: EFold1Inner"
   EIdx1{} -> error "Compile: not implemented: EIdx1"
 
+compileAssign :: String -> SList (Const String) env -> Ex env t -> CompM String
+compileAssign prefix env e = do
+  e' <- compile' env e
+  name <- genName' prefix
+  emit $ SVarDecl True (repSTy (typeOf e)) name e'
+  return name
+
 data Increment = Increment | Decrement
   deriving (Show)
 
@@ -1031,9 +1015,7 @@ compileScal pedantic typ x = case typ of
 compileExtremum :: String -> String -> String -> SList (Const String) env -> Ex env (TArr (S n) t) -> CompM CExpr
 compileExtremum nameBase opName operator env e = do
   let STArr (SS n) t = typeOf e
-  e' <- compile' env e
-  argname <- genName' (nameBase ++ "arg")
-  emit $ SVarDecl True (repSTy (STArr (SS n) t)) argname e'
+  argname <- compileAssign (nameBase ++ "arg") env e
 
   shszname <- genName' "shsz"
   -- This n is one less than the shape of the thing we're querying, which is