diff options
Diffstat (limited to 'src/CHAD/AST')
| -rw-r--r-- | src/CHAD/AST/Pretty.hs | 20 |
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 |