1 files changed, 11 insertions, 9 deletions
diff --git a/src/CHAD/AST/Pretty.hs b/src/CHAD/AST/Pretty.hs
index a9a8987..04c3d30 100644
--- a/src/CHAD/AST/Pretty.hs
+++ b/src/CHAD/AST/Pretty.hs
@@ -412,19 +412,19 @@ ppLam :: [ADoc] -> ADoc -> ADoc
 ppLam args body = ppString "(" <> hang 2 (ppString "\\" <> sep (args ++ [ppString "->"])
                                           <> softline <> body <> ppString ")")
 
-ppLoopNest :: SList f args -> LoopNest args outs -> String
+ppLoopNest :: SList f args -> LoopNest args outs bouts -> String
 ppLoopNest senv lnest = render $ fst . flip runM 1 $ do
   val <- mkSVal senv
   ppLoopNest' val lnest
 
 data RedKind = RKRet | RKBuild | RKSum
 
-ppLoopNest' :: SVal args -> LoopNest args outs -> M ADoc
+ppLoopNest' :: SVal args -> LoopNest args outs bouts -> M ADoc
 ppLoopNest' = \env lnest -> do
-  (f, outs) <- go env lnest
+  (f, outs, _bouts) <- go env lnest
   return (f (slistMap (\(Const _) -> Const RKRet) outs))
   where
-    go :: SVal args -> LoopNest args outs -> M (SList (Const RedKind) outs -> ADoc, SVal outs)
+    go :: SVal args -> LoopNest args outs bouts -> M (SList (Const RedKind) outs -> ADoc, SVal outs, SVal bouts)
     go env (Inner bs outs) = do
       (bs', names) <- goBindings env bs
       let outNames = slistMap (\i -> slistIdx (sappend names env) i) outs
@@ -434,19 +434,20 @@ ppLoopNest' = \env lnest -> do
       return (\kinds ->
                vcat (toList bs')
                <> hardline <> (annotate AKey (ppString "ret") <+> outDoc (unSList getConst kinds))
-             ,outNames)
+             ,outNames
+             ,SNil)
     go env (Layer bs1 wid lnest part bs2 outs) = do
       (bs1', names1) <- goBindings env bs1
       widname <- genName' "i"
-      (f, loopouts) <- go (Const widname `SCons` sappend names1 env) lnest
-      let (redkinds, mapouts, sumouts) = partition part loopouts
+      (f, loopouts, loopbouts) <- go (Const widname `SCons` sappend names1 env) lnest
+      let (redkinds, newbouts, sumouts) = partition part loopouts
       let lnest' = f redkinds
       (bs2', names2) <- goBindings (sappend sumouts (sappend names1 env)) bs2
       let outNames = slistMap (\i -> slistIdx (sappend names2 (sappend names1 env)) i) outs
           outDoc kinds =
             [annotate AKey (ppString "ret") <+>
               ppString "["
-              <> mconcat (map ppString (intersperse ", " (unSList _ (slistZip kinds outNames))))
+              <> mconcat (map ppString (intersperse ", " (unSList (\(Product.Pair (Const k) (Const n)) -> decorate k n) (slistZip kinds outNames))))
               -- <> ppString "] ++ ["
               -- <> mconcat (map ppString (intersperse ", " (unSList getConst mapouts)))
               <> ppString "]"]
@@ -456,7 +457,8 @@ ppLoopNest' = \env lnest -> do
                               <> hardline <> lnest')] ++
                      toList bs2' ++
                      outDoc kinds)
-             ,sappend outNames mapouts)
+             ,outNames
+             ,sappend newbouts loopbouts)
 
     decorate :: RedKind -> String -> String
     decorate RKRet name = name