diff options
Diffstat (limited to 'src/AST/Pretty.hs')
| -rw-r--r-- | src/AST/Pretty.hs | 34 |
1 files changed, 32 insertions, 2 deletions
diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs index 9018602..afa62c6 100644 --- a/src/AST/Pretty.hs +++ b/src/AST/Pretty.hs @@ -72,6 +72,7 @@ genNameIfUsedIn' prefix ty idx ex _ -> return "_" | otherwise = genName' prefix +-- TODO: let this return a type-tagged thing so that name environments are more typed than Const genNameIfUsedIn :: STy a -> Idx env a -> Expr x env t -> M String genNameIfUsedIn = \t -> genNameIfUsedIn' (nameBaseForType t) t @@ -209,8 +210,7 @@ ppExpr' d val expr = case expr of a' <- ppExpr' 0 (Const name2 `SCons` Const name1 `SCons` val) a b' <- ppExpr' 11 val b c' <- ppExpr' 11 val c - let opname = case cm of Commut -> "fold1i(C)" - Noncommut -> "fold1i" + let opname = "fold1i" ++ ppCommut cm return $ ppParen (d > 10) $ ppApp (annotate AHighlight (ppString opname) <> ppX expr) [ppLam [ppString name1, ppString name2] a', b', c'] @@ -235,6 +235,32 @@ ppExpr' d val expr = case expr of e' <- ppExpr' 11 val e return $ ppParen (d > 10) $ ppString "minimum1i" <> ppX expr <+> e' + EFold1InnerD1 _ cm a b c -> do + name1 <- genNameIfUsedIn (typeOf b) (IS IZ) a + name2 <- genNameIfUsedIn (typeOf b) IZ a + a' <- ppExpr' 0 (Const name2 `SCons` Const name1 `SCons` val) a + b' <- ppExpr' 11 val b + c' <- ppExpr' 11 val c + let opname = "fold1iD1" ++ ppCommut cm + return $ ppParen (d > 10) $ + ppApp (annotate AHighlight (ppString opname) <> ppX expr) [ppLam [ppString name1, ppString name2] a', b', c'] + + EFold1InnerD2 _ cm t2 ef ep ezi ebog ed -> do + let STArr _ (STPair t1 ttape) = typeOf ebog + name1 <- genNameIfUsedIn ttape (IS (IS (IS IZ))) ef + name2 <- genNameIfUsedIn t1 (IS (IS IZ)) ef + name3 <- genNameIfUsedIn t1 (IS IZ) ef + name4 <- genNameIfUsedIn (fromSMTy t2) IZ ef + ef' <- ppExpr' 0 (Const name4 `SCons` Const name3 `SCons` Const name2 `SCons` Const name1 `SCons` val) ef + ep' <- ppExpr' 11 val ep + ezi' <- ppExpr' 11 val ezi + ebog' <- ppExpr' 11 val ebog + ed' <- ppExpr' 11 val ed + let opname = "fold1iD2" ++ ppCommut cm + return $ ppParen (d > 10) $ + ppApp (annotate AHighlight (ppString opname) <> ppX expr) + [ppLam [ppString name1, ppString name2, ppString name3, ppString name4] ef', ep', ezi', ebog', ed'] + EConst _ ty v | Dict <- scalRepIsShow ty -> return $ ppString (showsPrec d v "") <> ppX expr @@ -386,6 +412,10 @@ ppSparse (SMTMaybe t) (SpMaybe s) = "M" ++ ppSparse t s ppSparse (SMTArr _ t) (SpArr s) = "A" ++ ppSparse t s ppSparse (SMTScal _) SpScal = "." +ppCommut :: Commutative -> String +ppCommut Commut = "(C)" +ppCommut Noncommut = "" + ppX :: PrettyX x => Expr x env t -> ADoc ppX expr = annotate AExt $ ppString $ prettyXsuffix (extOf expr) |
