Debug printing of expression tree (default off)

1 files changed, 41 insertions, 22 deletions

diff --git a/SC/Exp.hs b/SC/Exp.hs
index ef18a7f..cf4e096 100644
--- a/SC/Exp.hs
+++ b/SC/Exp.hs
@@ -4,12 +4,16 @@ module SC.Exp where
 
 import qualified Data.Array.Accelerate.AST as A
 import Data.Array.Accelerate.AST.LeftHandSide
+import Data.Array.Accelerate.AST.Idx (idxToInt)
 import Data.Array.Accelerate.AST.Var
 import Data.Array.Accelerate.Representation.Array
 import Data.Array.Accelerate.Representation.Shape
 import Data.Array.Accelerate.Representation.Type
 import Data.Array.Accelerate.Type
 
+-- import Debug.Trace
+
+import Debug
 import qualified Language.C as C
 import SC.Defs
 import SC.Monad
@@ -41,7 +45,8 @@ compileFun aenv (A.Lam lhs (A.Body body)) = do
     (argnames, env) <- genVarsEnv lhs VENil
     outnames <- itupmap (\(TypedName t n) -> TypedName (C.TPtr t) n)
                     <$> genVars (A.expType body)
-    (usedA, res) <- compileExp' aenv env body
+    ((_tree, usedA), res) <- compileExp' aenv env body
+    -- traceM ("Compiled expression:\n" ++ prettyTree "  " "  " tree)
     (sts1, retexprs) <- toStExprs (A.expType body) res
     let sts2 = genoutstores outnames retexprs
         arrayarguments =
@@ -71,27 +76,38 @@ compileFun _ _ = error "compileFun: Not single-argument function"
 compileExp :: AVarEnv aenv -> A.Exp aenv t -> SC (CompiledFun aenv () t)
 compileExp aenv expr = compileFun aenv (A.Lam (LeftHandSideWildcard TupRunit) (A.Body expr))
 
+data Tree = Node String [Tree] | Leaf String
+
+prettyTree :: String -> String -> Tree -> String
+prettyTree pre _    (Leaf s) = pre ++ s ++ "\n"
+prettyTree pre pre2 (Node s []) = prettyTree pre pre2 (Leaf s)
+prettyTree pre pre2 (Node s ts) =
+    let (ts1, t2) = (init ts, last ts)
+    in pre ++ s ++ "\n" ++ concatMap (prettyTree (pre2 ++ "├─") (pre2 ++ "│ ")) ts1 ++ prettyTree (pre2 ++ "└─") (pre2 ++ "  ") t2
+
 compileExp' :: AVarEnv aenv -> VarEnv env -> A.OpenExp env aenv t
-            -> SC ([SomeArray], Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
+            -> SC ((Tree, [SomeArray]), Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
 compileExp' aenv env = \case
+    -- Foreign, IndexSlice, IndexFull, FromIndex, Case, Cond, While, PrimConst, ShapeSize, Undef, Coerce
+
     A.Let lhs rhs body -> do
         (names, env') <- genVarsEnv lhs env
         let sts1 = [C.SDecl t n Nothing | TypedName t n <- itupList names]
-        (usedA2, sts2) <- fmap (`toStoring` names) <$> compileExp' aenv env rhs
-        (usedA3, res3) <- compileExp' aenv env' body
-        return (usedA2 ++ usedA3
+        ((tree2, usedA2), sts2) <- fmap (`toStoring` names) <$> compileExp' aenv env rhs
+        ((tree3, usedA3), res3) <- compileExp' aenv env' body
+        return ((Node ("Let [" ++ show (length (itupList names)) ++ " vars]") [tree2, tree3], usedA2 ++ usedA3)
                ,fmap (\(sts, exs) -> (sts1 ++ sts2 ++ sts, exs)) res3)
 
     A.Evar (Var _ idx) ->
-        return ([], Right ([], ITupSingle (C.EVar (veprj env idx))))
+        return ((Leaf ("Evar " ++ show (idxToInt idx)), []), Right ([], ITupSingle (C.EVar (veprj env idx))))
 
     A.Nil ->
-        return ([], Right ([], ITupIgnore))
+        return ((Leaf "Nil", []), Right ([], ITupIgnore))
 
     A.Pair a b -> do
-        (usedA1, res1) <- compileExp' aenv env a
-        (usedA2, res2) <- compileExp' aenv env b
-        return (usedA1 ++ usedA2
+        ((tree1, usedA1), res1) <- compileExp' aenv env a
+        ((tree2, usedA2), res2) <- compileExp' aenv env b
+        return ((Node "Pair" [tree1, tree2], usedA1 ++ usedA2)
                ,case (res1, res2) of
                   (Right (sts1, exp1), Right (sts2, exp2)) ->
                       Right (sts1 ++ sts2, ITupPair exp1 exp2)
@@ -102,20 +118,21 @@ compileExp' aenv env = \case
 
     A.Const ty x
       | Just str <- showExpConst ty x
-      -> return ([], Right ([], ITupSingle (C.ELit str)))
+      -> return ((Leaf ("Const (" ++ str ++ ")"), []), Right ([], ITupSingle (C.ELit str)))
 
     A.PrimApp (A.PrimAdd _) e -> binary aenv env "+" e
     A.PrimApp (A.PrimSub _) e -> binary aenv env "-" e
     A.PrimApp (A.PrimMul _) e -> binary aenv env "*" e
     A.PrimApp (A.PrimQuot _) e -> binary aenv env "/" e
     A.PrimApp (A.PrimRem _) e -> binary aenv env "%" e
+    A.PrimApp op _ -> throw $ "Unsupported Exp primitive operator: " ++ showPrimFun op
 
     A.Shape (Var _ idx) ->
         let (shnames, _) = aveprj aenv idx
             buildExprs :: ShNames sh -> Exprs sh
             buildExprs ShZ = ITupIgnore
             buildExprs (ShS names n) = ITupPair (buildExprs names) (ITupSingle (C.EVar n))
-        in return ([], Right ([], buildExprs shnames))
+        in return ((Leaf ("Shape a" ++ show (idxToInt idx)), []), Right ([], buildExprs shnames))
 
     A.ToIndex shr she idxe -> do
         let build :: ShapeR sh -> Exprs sh -> Exprs sh -> C.Expr
@@ -125,11 +142,11 @@ compileExp' aenv env = \case
                                     (ITupPair idxes' (ITupSingle idxe')) =
                 C.EOp (C.EOp (build shr' shes' idxes') "*" she') "+" idxe'
             build _ _ _ = error "wat"
-        (usedA1, res1) <- compileExp' aenv env she
+        ((tree1, usedA1), res1) <- compileExp' aenv env she
         (sts1, shes) <- toStExprs (shapeType shr) res1
-        (usedA2, res2) <- compileExp' aenv env idxe
+        ((tree2, usedA2), res2) <- compileExp' aenv env idxe
         (sts2, idxes) <- toStExprs (shapeType shr) res2
-        return (usedA1 ++ usedA2, Right (sts1 ++ sts2, ITupSingle (build shr shes idxes)))
+        return ((Node "ToIndex" [tree1, tree2], usedA1 ++ usedA2), Right (sts1 ++ sts2, ITupSingle (build shr shes idxes)))
 
     A.Index avar@(Var (ArrayR shr _) _) she ->
         compileExp' aenv env $
@@ -138,22 +155,24 @@ compileExp' aenv env = \case
     A.LinearIndex (Var _ idx) e -> do
         temp <- genName "i"
         let sts0 = [C.SDecl (C.TInt C.B64) temp Nothing]
-        (usedA1, sts1) <- fmap (`toStoring` ITupSingle (TypedName (C.TInt C.B64) temp))
-                              <$> compileExp' aenv env e
+        ((tree1, usedA1), sts1) <-
+            fmap (`toStoring` ITupSingle (TypedName (C.TInt C.B64) temp))
+              <$> compileExp' aenv env e
         let (shnames, anames) = aveprj aenv idx
             usedA = SomeArray shnames anames : usedA1
-        return (usedA, Right (sts0 ++ sts1
-                             ,itupmap (\(TypedAName _ name) -> C.EIndex name (C.EVar temp)) anames))
+        return ((Node ("LinearIndex a" ++ show (idxToInt idx)) [tree1], usedA)
+               ,Right (sts0 ++ sts1
+                      ,itupmap (\(TypedAName _ name) -> C.EIndex name (C.EVar temp)) anames))
 
     e -> throw $ "Unsupported Exp constructor: " ++ A.showExpOp e
   where
     binary :: AVarEnv aenv -> VarEnv env -> String -> A.OpenExp env aenv (a, b)
-           -> SC ([SomeArray], Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
+           -> SC ((Tree, [SomeArray]), Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
     binary aenv' env' op e' = do
-        (usedA, res) <- compileExp' aenv' env' e'
+        ((tree, usedA), res) <- compileExp' aenv' env' e'
         (sts, ITupPair (ITupSingle e1) (ITupSingle e2)) <-
             toStExprs (A.expType e') res
-        return (usedA, Right (sts, ITupSingle (C.EOp e1 op e2)))
+        return ((Node ("binary " ++ show op) [tree], usedA), Right (sts, ITupSingle (C.EOp e1 op e2)))
 
 toStExprs :: TypeR t -> Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t) -> SC ([C.Stmt], Exprs t)
 toStExprs ty (Left fun) = do