1 files changed, 29 insertions, 33 deletions

diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index dbbc021..5610d36 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -1,16 +1,15 @@
-{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE DataKinds #-}
-{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE EmptyCase #-}
 {-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE PolyKinds #-}
-{-# LANGUAGE EmptyCase #-}
-{-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE TupleSections #-}
-module AST.Pretty where
+{-# LANGUAGE TypeOperators #-}
+module AST.Pretty (ppExpr) where
 
 import Control.Monad (ap)
 import Data.List (intersperse)
-import Data.Foldable (toList)
 import Data.Functor.Const
 
 import AST
@@ -29,10 +28,6 @@ valprj (VPush x _) IZ = x
 valprj (VPush _ env) (IS i) = valprj env i
 valprj VTop i = case i of {}
 
-vpushN :: Vec n a -> Val (Const a) env -> Val (Const a) (ConsN n TIx env)
-vpushN VNil v = v
-vpushN (name :< names) v = VPush (Const name) (vpushN names v)
-
 newtype M a = M { runM :: Int -> (a, Int) }
   deriving (Functor)
 instance Applicative M where { pure x = M (\i -> (x, i)) ; (<*>) = ap }
@@ -115,12 +110,10 @@ ppExpr' d val = \case
 
   EBuild _ n a b -> do
     a' <- ppExpr' 11 val a
-    names <- sequence (vecGenerate n (\_ -> genName))  -- TODO generate underscores
-    e' <- ppExpr' 0 (vpushN names val) b
+    name <- genNameIfUsedIn (tTup (sreplicate n tIx)) IZ b
+    e' <- ppExpr' 0 (VPush (Const name) val) b
     return $ showParen (d > 10) $
-      showString "build " . a' . showString " (\\["
-      . foldr (.) id (intersperse (showString ",") (map showString (reverse (toList names))))
-      . showString ("] -> ") . e' . showString ")"
+      showString "build " . a' . showString (" (\\" ++ name ++ " -> ") . e' . showString ")"
 
   EFold1 _ a b -> do
     name1 <- genNameIfUsedIn (typeOf a) (IS IZ) a
@@ -135,9 +128,9 @@ ppExpr' d val = \case
     e' <- ppExpr' 11 val e
     return $ showParen (d > 10) $ showString "unit " . e'
 
-  EReplicate _ e -> do
-    e' <- ppExpr' 11 val e
-    return $ showParen (d > 10) $ showString "replicate " . e'
+  -- EReplicate _ e -> do
+  --   e' <- ppExpr' 11 val e
+  --   return $ showParen (d > 10) $ showString "replicate " . e'
 
   EConst _ ty v -> return $ showString $ case ty of
     STI32 -> show v ; STI64 -> show v ; STF32 -> show v ; STF64 -> show v ; STBool -> show v
@@ -151,14 +144,15 @@ ppExpr' d val = \case
     b' <- ppExpr' 9 val b
     return $ showParen (d > 8) $ a' . showString " ! " . b'
 
-  EIdx _ e es -> do
-    e' <- ppExpr' 9 val e
-    es' <- traverse (ppExpr' 0 val) es
+  EIdx _ _ a b -> do
+    a' <- ppExpr' 9 val a
+    b' <- ppExpr' 10 val b
     return $ showParen (d > 8) $
-      e' . showString " ! "
-      . showString "["
-      . foldr (.) id (intersperse (showString ", ") (reverse (toList es')))
-      . showString "]"
+      a' . showString " !! " . b'
+
+  EShape _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ showParen (d > 10) $ showString "shape " . e'
 
   EOp _ op (EPair _ a b)
     | (Infix, ops) <- operator op -> do
@@ -175,30 +169,30 @@ ppExpr' d val = \case
 
   EWith e1 e2 -> do
     e1' <- ppExpr' 11 val e1
-    let STArr n t = typeOf e1
-    name <- genNameIfUsedIn' "ac" (STAccum n t) IZ e2
-    e2' <- ppExpr' 11 (VPush (Const name) val) e2
+    name <- genNameIfUsedIn' "ac" (STAccum (typeOf e1)) IZ e2
+    e2' <- ppExpr' 0 (VPush (Const name) val) e2
     return $ showParen (d > 10) $
       showString "with " . e1' . showString (" (\\" ++ name ++ " -> ")
       . e2' . showString ")"
 
-  EAccum1 e1 e2 e3 -> do
+  EAccum i e1 e2 e3 -> do
     e1' <- ppExpr' 11 val e1
     e2' <- ppExpr' 11 val e2
     e3' <- ppExpr' 11 val e3
     return $ showParen (d > 10) $
-      showString "accum1 " . e1' . showString " " . e2' . showString " " . e3'
+      showString ("accum " ++ show (unSNat i) ++ " ") . e1' . showString " " . e2' . showString " " . e3'
 
   EError _ s -> return $ showParen (d > 10) $ showString ("error " ++ show s)
 
 ppExprLet :: Int -> SVal env -> Expr x env t -> M ShowS
 ppExprLet d val etop = do
-  let collect :: SVal env -> Expr x env t -> M ([(String, ShowS)], ShowS)
+  let collect :: SVal env -> Expr x env t -> M ([(String, Occ, ShowS)], ShowS)
       collect val' (ELet _ rhs body) = do
+        let occ = occCount IZ body
         name <- genNameIfUsedIn (typeOf rhs) IZ body
         rhs' <- ppExpr' 0 val' rhs
         (binds, core) <- collect (VPush (Const name) val') body
-        return ((name, rhs') : binds, core)
+        return ((name, occ, rhs') : binds, core)
       collect val' e = ([],) <$> ppExpr' 0 val' e
 
   (binds, core) <- collect val etop
@@ -210,7 +204,9 @@ ppExprLet d val etop = do
     showString ("let " ++ open)
     . foldr (.) id
         (intersperse (showString " ; ")
-           (map (\(name, rhs) -> showString (name ++ " = ") . rhs) binds))
+           (map (\(name, _occ, rhs) ->
+                    showString (name ++ {- " (" ++ show _occ ++ ")" ++ -} " = ") . rhs)
+                binds))
     . showString (close ++ " in ")
     . core