1 files changed, 242 insertions, 38 deletions

diff --git a/Simplify.hs b/Simplify.hs
index 9ceaef9..e95043d 100644
--- a/Simplify.hs
+++ b/Simplify.hs
@@ -2,12 +2,15 @@
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
 module Simplify (
     simplify,
     simplifyFix,
+    simbeta, simpair, simindex, simifold1,
+    simfix, SimList(..),
 ) where
 
 import Data.Bifunctor
@@ -16,26 +19,51 @@ import qualified Data.Kind as Kind
 import Data.List (find)
 import Data.Type.Equality
 
+import Debug.Trace
+
 import AST
+import Count
 import Sink
 
 
+data Bound a = Inclusive a | Exclusive a
+  deriving (Show)
+
+instance Functor Bound where
+    fmap f (Inclusive x) = Inclusive (f x)
+    fmap f (Exclusive x) = Exclusive (f x)
+
 data family Info (env :: [Kind.Type]) a
--- data instance Info env Int = InfoInt
--- data instance Info env Bool = InfoBool
--- data instance Info env Double = InfoDouble
+-- | Lower bound (inclusive), upper bound (inclusive/exclusive)
+data instance Info env Int = InfoInt (Maybe (Exp env Int)) (Maybe (Bound (Exp env Int)))
 data instance Info env (Array sh t) = InfoArray (Exp env sh)
--- data instance Info env () = InfoNil
-data instance Info env (a, b) = InfoPair (Info env a) (Info env b)
--- data instance Info env (a -> b) = InfoFun
+data instance Info env () = InfoNil
+data instance Info env (a, b) = InfoPair (Maybe (Info env a)) (Maybe (Info env b))
 
 data IEnv env where
     ITop :: IEnv env
     ICons :: Type a -> Maybe (Info (a ': env) a) -> IEnv env -> IEnv (a ': env)
 
+showsInfo :: Int -> Type a -> Info env a -> ShowS
+showsInfo d TInt (InfoInt a b) = showParen (d > 10) $
+    showString "InfoInt " . showsPrec 11 a . showString " " . showsPrec 11 b
+showsInfo d TArray{} (InfoArray a) = showParen (d > 10) $
+    showString "InfoArray " . showsPrec 11 a
+showsInfo _ TNil InfoNil = showString "InfoNil"
+showsInfo d (TPair t1 t2) (InfoPair a b) = showParen (d > 10) $
+    showString "InfoPair " . showsInfo' 11 t1 a . showString " " . showsInfo' 11 t2 b
+showsInfo _ _ _ = error "showsInfo: No definition"
+
+showsInfo' :: Int -> Type a -> Maybe (Info env a) -> ShowS
+showsInfo' _ _ Nothing = showString "Nothing"
+showsInfo' d t (Just x) = showParen (d > 10) $
+    showString "Just " . showsInfo 11 t x
+
 sinkInfo1 :: Type a -> Info env a -> Info (t ': env) a
+sinkInfo1 TInt (InfoInt a b) = InfoInt (sinkExp1 <$> a) (fmap sinkExp1 <$> b)
 sinkInfo1 TArray{} (InfoArray e) = InfoArray (sinkExp1 e)
-sinkInfo1 (TPair t1 t2) (InfoPair a b) = InfoPair (sinkInfo1 t1 a) (sinkInfo1 t2 b)
+sinkInfo1 TNil InfoNil = InfoNil
+sinkInfo1 (TPair t1 t2) (InfoPair a b) = InfoPair (sinkInfo1 t1 <$> a) (sinkInfo1 t2 <$> b)
 sinkInfo1 _ _ = error "Unknown info in sinkInfo1"
 
 iprj :: IEnv env -> Idx env a -> Maybe (Type a, Info env a)
@@ -64,6 +92,8 @@ simplify' env = \case
         let (arg', info) = simplify' env arg
             env' = ICons (typeof arg) (sinkInfo1 (typeof arg) <$> info) env
         in (simplifyLet arg' (fst (simplify' env' e)), Nothing)
+    Lit (LInt n) -> (Lit (LInt n), Just (InfoInt (Just (Lit (LInt n)))
+                                                 (Just (Inclusive (Lit (LInt n))))))
     Lit l -> (Lit l, Nothing)
     Cond a b c ->
         (Cond (fst (simplify' env a)) (fst (simplify' env b)) (fst (simplify' env c)), Nothing)
@@ -71,27 +101,55 @@ simplify' env = \case
     Pair a b ->
         let (a', ia) = simplify' env a
             (b', ib) = simplify' env b
-        in (Pair a' b', InfoPair <$> ia <*> ib)
-    Fst e -> bimap simplifyFst (fmap (\(InfoPair i _) -> i)) (simplify' env e)
-    Snd e -> bimap simplifySnd (fmap (\(InfoPair _ i) -> i)) (simplify' env e)
+        in (simplifyPair a' b', Just (InfoPair ia ib))
+    Fst e -> bimap simplifyFst (>>= (\(InfoPair i _) -> i)) (simplify' env e)
+    Snd e -> bimap simplifySnd (>>= (\(InfoPair _ i) -> i)) (simplify' env e)
+    Build sht a (Lam shty fe) ->
+        let a' = fst (simplify' env a)
+            env' = ICons shty (Just (shapeBoundInfo sht (sinkExp1 a'))) env
+        in (Build sht a' (Lam shty (fst (simplify' env' fe))), Just (InfoArray a'))
     Build sht a b ->
         let a' = fst (simplify' env a)
         in (Build sht a' (fst (simplify' env b)), Just (InfoArray a'))
-    Ifold sht a b c -> (Ifold sht (fst (simplify' env a)) (fst (simplify' env b)) (fst (simplify' env c)), Nothing)
+    Ifold sht a b c ->
+        (simplifyIfold env sht (fst (simplify' env a)) (fst (simplify' env b)) (fst (simplify' env c)), Nothing)
     Index a b -> (simplifyIndex (fst (simplify' env a)) (fst (simplify' env b)), Nothing)
     Shape e ->
         case simplify' env e of
           (_, Just (InfoArray she)) -> (she, Nothing)
           (e', _) -> (Shape e', Nothing)
+    Undef t -> (Undef t, Nothing)
+
+shapeBoundInfo :: ShapeType sh -> Exp env sh -> Info env sh
+shapeBoundInfo STZ _ = InfoNil
+shapeBoundInfo (STC sht) she =
+    InfoPair (Just (shapeBoundInfo sht (Fst she)))
+             (Just (InfoInt (Just (Lit (LInt 0))) (Just (Exclusive (Snd she)))))
 
 simplifyApp :: Exp env (a -> b) -> Exp env a -> Exp env b
 simplifyApp (Const CAddI) (Pair (Lit (LInt a)) (Lit (LInt b))) = Lit (LInt (a + b))
+simplifyApp (Const CAddI) (Pair a (Lit (LInt 0))) = a
+simplifyApp (Const CAddI) (Pair (Lit (LInt 0)) a) = a
+-- simplifyApp (Const CAddI) (Pair a b) | Just Refl <- geq a b =
+--     simplifyApp (Const CMulI) (Pair (Lit (LInt 2)) a)
 simplifyApp (Const CSubI) (Pair (Lit (LInt a)) (Lit (LInt b))) = Lit (LInt (a - b))
 simplifyApp (Const CMulI) (Pair (Lit (LInt a)) (Lit (LInt b))) = Lit (LInt (a * b))
+simplifyApp (Const CMulI) (Pair a (Lit (LInt 1))) = a
+simplifyApp (Const CMulI) (Pair (Lit (LInt 1)) a) = a
+simplifyApp (Const CMulI) (Pair _ (Lit (LInt 0))) = Lit (LInt 0)
+simplifyApp (Const CMulI) (Pair (Lit (LInt 0)) _) = Lit (LInt 0)
 simplifyApp (Const CDivI) (Pair (Lit (LInt a)) (Lit (LInt b))) = Lit (LInt (a `div` b))
 simplifyApp (Const CAddF) (Pair (Lit (LDouble a)) (Lit (LDouble b))) = Lit (LDouble (a + b))
+simplifyApp (Const CAddF) (Pair a (Lit (LDouble 0))) = a
+simplifyApp (Const CAddF) (Pair (Lit (LDouble 0)) a) = a
+-- simplifyApp (Const CAddF) (Pair a b) | Just Refl <- geq a b =
+--     simplifyApp (Const CMulF) (Pair (Lit (LDouble 2)) a)
 simplifyApp (Const CSubF) (Pair (Lit (LDouble a)) (Lit (LDouble b))) = Lit (LDouble (a - b))
 simplifyApp (Const CMulF) (Pair (Lit (LDouble a)) (Lit (LDouble b))) = Lit (LDouble (a * b))
+simplifyApp (Const CMulF) (Pair a (Lit (LDouble 1))) = a
+simplifyApp (Const CMulF) (Pair (Lit (LDouble 1)) a) = a
+simplifyApp (Const CMulF) (Pair _ (Lit (LDouble 0))) = Lit (LDouble 0)
+simplifyApp (Const CMulF) (Pair (Lit (LDouble 0)) _) = Lit (LDouble 0)
 simplifyApp (Const CDivF) (Pair (Lit (LDouble a)) (Lit (LDouble b))) = Lit (LDouble (a / b))
 simplifyApp (Const CLog) (Lit (LDouble a)) = Lit (LDouble (log a))
 simplifyApp (Const CExp) (Lit (LDouble a)) = Lit (LDouble (exp a))
@@ -107,15 +165,17 @@ simplifyApp (Const COr) (Pair (Lit (LBool a)) (Lit (LBool b))) = Lit (LBool (a |
 simplifyApp (Const CNot) (Lit (LBool a)) = Lit (LBool (not a))
 
 simplifyApp (Lam _ e) arg
-  | isDuplicable arg || countOcc Zero e <= 1
+  | isDuplicable arg || usesOf Zero e <= 1
   = simplify (subst arg e)
 simplifyApp (Lam _ e) arg = simplifyLet arg e
 
+simplifyApp f (Cond c a b) = simplify $ Cond c (App f a) (App f b)
+
 simplifyApp a b = App a b
 
 simplifyLet :: Exp env a -> Exp (a ': env) b -> Exp env b
 simplifyLet arg e
-  | isDuplicable arg || countOcc Zero e <= 1
+  | isDuplicable arg || usesOf Zero e <= 1
   = simplify (subst arg e)
 simplifyLet (Pair a b) e =
     simplifyLet a $
@@ -124,24 +184,89 @@ simplifyLet (Pair a b) e =
                                              (Var (typeof b) Zero)
                                 Succ i -> Var t (Succ (Succ i)))
                    e
-simplifyLet (Cond c a b) e
-  | isDuplicable a && isDuplicable b
-  = simplifyLet c $
-      (subst' (\t -> \case Zero -> Cond (Var TBool Zero) (sinkExp1 a) (sinkExp1 b)
-                           Succ i -> Var t (Succ i))
-              e)
-simplifyLet a b = Let (simplify a) (simplify b)
+-- simplifyLet (Cond c a b) e
+--   | isDuplicable a && isDuplicable b
+--   = simplifyLet c $
+--       (subst' (\t -> \case Zero -> Cond (Var TBool Zero) (sinkExp1 a) (sinkExp1 b)
+--                            Succ i -> Var t (Succ i))
+--               e)
+simplifyLet (Cond c a b) e = simplify $ Cond c (Let a e) (Let b e)
+simplifyLet a b = Let a b
+
+simplifyPair :: Exp env a -> Exp env b -> Exp env (a, b)
+simplifyPair (Cond c a b) d = simplify $ Cond c (Pair a d) (Pair b d)
+simplifyPair d (Cond c a b) = simplify $ Cond c (Pair d a) (Pair d b)
+simplifyPair a b = Pair a b
 
 simplifyFst :: Exp env (a, b) -> Exp env a
 simplifyFst (Pair e _) = e
 simplifyFst (Let a e) = simplifyLet a (simplifyFst e)
+simplifyFst (Cond c a b) = simplify $ Cond c (Fst a) (Fst b)
 simplifyFst e = Fst e
 
 simplifySnd :: Exp env (a, b) -> Exp env b
 simplifySnd (Pair _ e) = e
 simplifySnd (Let a e) = simplifyLet a (simplifySnd e)
+simplifySnd (Cond c a b) = simplify $ Cond c (Snd a) (Snd b)
 simplifySnd e = Snd e
 
+simplifyIfold :: IEnv env -> ShapeType sh -> Exp env ((a, sh) -> a) -> Exp env a -> Exp env sh -> Exp env a
+simplifyIfold env sht fe e0 she
+  | Just res <- splitIfold sht fe e0 she
+  = fst (simplify' env res)
+-- Given the following:
+--   ifold (\(a,i) -> if i == cmpref then val else a) _ she
+-- and given that we can prove that 0 <= cmpref < she and that 'val' is free,
+-- the whole fold can be replaced with 'val'.
+simplifyIfold env sht (Lam argty (Cond (App (Const (CEq _)) (Pair (Snd (Var _ Zero)) cmpref)) val (Fst (Var _ Zero)))) e0 she
+  | let env' = ICons argty (Just (InfoPair Nothing (Just (shapeBoundInfo sht (sinkExp1 she))))) env
+  , trace ("si: trying") True
+  , proveShapeBound env' CLeI sht (zeroShapeExp sht) cmpref
+  , trace ("si: prf1 = True") True
+  , proveShapeBound env' CLtI sht cmpref (sinkExp1 she)
+  , trace ("si: prf2 = True") True
+  , trace ("si: cmpref = " ++ show val) True
+  , usesOf Zero cmpref == 0
+  = simplifyLet (Pair e0 (subst (error "usesOf == 0 was wrong") cmpref)) val
+simplifyIfold _ sht fe e0 she = Ifold sht fe e0 she
+
+zeroShapeExp :: ShapeType sh -> Exp env sh
+zeroShapeExp STZ = Lit LNil
+zeroShapeExp (STC sht) = Pair (zeroShapeExp sht) (Lit (LInt 0))
+
+proveShapeBound :: IEnv env -> Constant ((Int, Int) -> Bool) -> ShapeType sh -> Exp env sh -> Exp env sh -> Bool
+proveShapeBound _ _ STZ _ _ = True
+proveShapeBound env cmpop (STC sht) e1 e2 =
+    let (_, info1) = simplify' env (Snd e1)
+        (_, info2) = simplify' env (Snd e2)
+        inclLo2 = case info2 of
+                    Just (InfoInt (Just lo2) _) -> lo2
+                    _ -> Snd e2  -- this is also an inclusive lower bound, after all
+        restresult = proveShapeBound env cmpop sht (Fst e1) (Fst e2)
+    in restresult && case (cmpop, info1) of
+         (CLeI, Just (InfoInt _ (Just (Inclusive hi1)))) ->
+             proveLe hi1 inclLo2
+         (CLtI, Just (InfoInt _ (Just (Exclusive hi1)))) ->
+             proveLe hi1 inclLo2
+         _ -> trace ("proveShapeBound: " ++ show cmpop ++ " (" ++ show e1 ++ ") (" ++ show e2 ++ ")") $
+              trace ("  e1 = " ++ show e1) $
+              trace ("  e2 = " ++ show e2) $
+              trace ("  info1 = " ++ showsInfo' 0 TInt info1 "") $
+              trace ("  info2 = " ++ showsInfo' 0 TInt info2 "") $
+              trace ("  inclLo2 = " ++ show inclLo2) $
+              False
+
+proveLe :: Exp env Int -> Exp env Int -> Bool
+proveLe = \e1 e2 ->
+    let res = proveLe' e1 e2
+    in trace ("proveLe: '" ++ show e1 ++ "'  <=  '" ++ show e2 ++ "'  ->  " ++ show res)
+             res
+
+proveLe' :: Exp env Int -> Exp env Int -> Bool
+proveLe' e1 e2 | Just Refl <- geq e1 e2 = True
+proveLe' (Lit (LInt a)) (Lit (LInt b)) | a <= b = True
+proveLe' _ _ = False
+
 simplifyIndex :: Exp env (Array sh a) -> Exp env sh -> Exp env a
 simplifyIndex (Build _ _ f) e = simplifyApp f e
 simplifyIndex a e = Index a e
@@ -162,24 +287,6 @@ isDuplicable (Fst e) = isDuplicable e
 isDuplicable (Snd e) = isDuplicable e
 isDuplicable _ = False
 
-countOcc :: Idx env t -> Exp env a -> Int
-countOcc i (App a b) = countOcc i a + countOcc i b
-countOcc i (Lam _ e) = countOcc (Succ i) e
-countOcc i (Var _ j)
-  | Just Refl <- geq i j = 1
-  | otherwise = 0
-countOcc i (Let a b) = countOcc i a + countOcc (Succ i) b
-countOcc _ (Lit _) = 0
-countOcc i (Cond a b c) = countOcc i a + countOcc i b + countOcc i c
-countOcc _ (Const _) = 0
-countOcc i (Pair a b) = countOcc i a + countOcc i b
-countOcc i (Fst e) = countOcc i e
-countOcc i (Snd e) = countOcc i e
-countOcc i (Build _ a b) = countOcc i a + countOcc i b
-countOcc i (Ifold _ a b c) = countOcc i a + countOcc i b + countOcc i c
-countOcc i (Index a b) = countOcc i a + countOcc i b
-countOcc i (Shape e) = countOcc i e
-
 subst :: Exp env t -> Exp (t ': env) a -> Exp env a
 subst arg e = subst' (\t -> \case Zero -> arg ; Succ i -> Var t i) e
 
@@ -201,3 +308,100 @@ subst' f (Build sht a b) = Build sht (subst' f a) (subst' f b)
 subst' f (Ifold sht a b c) = Ifold sht (subst' f a) (subst' f b) (subst' f c)
 subst' f (Index a b) = Index (subst' f a) (subst' f b)
 subst' f (Shape e) = Shape (subst' f e)
+subst' _ (Undef t) = Undef t
+
+splitIfold :: ShapeType sh -> Exp env ((s, sh) -> s) -> Exp env s -> Exp env sh -> Maybe (Exp env s)
+splitIfold sht (Lam (TPair (TPair t1 t2) tidx) (Pair e1 e2)) e0 she
+  | let uses1 = usesOf' PathStart Zero e1
+        uses2 = usesOf' PathStart Zero e2
+  , lycontract (lysnd (lyfst uses1)) == 0
+  , lycontract (lyfst (lyfst uses2)) == 0
+    -- Substitute the argument in e1 and t2 to refer to just the used
+    -- components of their argument. To do this we reconstruct the original,
+    -- partially unused argument by putting an 'Undef' in the unused spot.
+  , let e1' = subst' (\t -> \case Zero ->
+                                    Pair (Pair (Fst (Var (TPair t1 tidx) Zero))
+                                               (Undef t2))
+                                         (Snd (Var (TPair t1 tidx) Zero))
+                                  Succ i -> Var t (Succ i))
+                     e1
+        e2' = subst' (\t -> \case Zero ->
+                                    Pair (Pair (Undef t1)
+                                               (Fst (Var (TPair t2 tidx) Zero)))
+                                         (Snd (Var (TPair t2 tidx) Zero))
+                                  Succ i -> Var t (Succ i))
+                     e2
+  = Just $
+      Let e0 $ Let (sinkExp1 she) $
+        Pair (Ifold sht (sinkExp2 (Lam (TPair t1 tidx) e1'))
+                        (Fst (Var (TPair t1 t2) (Succ Zero)))
+                        (Var tidx Zero))
+             (Ifold sht (sinkExp2 (Lam (TPair t2 tidx) e2'))
+                        (Snd (Var (TPair t1 t2) (Succ Zero)))
+                        (Var tidx Zero))
+splitIfold _ _ _ _ = Nothing
+
+simbeta :: Exp env a -> Exp env a
+simbeta = \case
+    App (Lam _ e) a
+      | isDuplicable a || usesOf Zero e <= 1
+      -> simbeta (subst a e)
+      | otherwise
+      -> Let (simbeta a) (simbeta e)
+    Let a e
+      | isDuplicable a || usesOf Zero e <= 1
+      -> simbeta (subst a e)
+    e -> simrecurse simbeta e
+
+simpair :: Exp env a -> Exp env a
+simpair = \case
+    Fst (Pair a _) -> simpair a
+    Snd (Pair _ b) -> simpair b
+    Fst (Let a b) -> Let (simpair a) (simpair (Fst b))
+    Snd (Let a b) -> Let (simpair a) (simpair (Snd b))
+    e -> simrecurse simpair e
+
+simindex :: Exp env a -> Exp env a
+simindex = \case
+    Index (Build _ _ f) e ->
+        App (simindex f) (simindex e)
+    e -> simrecurse simindex e
+
+simifold1 :: Exp env a -> Exp env a
+simifold1 = \case
+    Ifold sht fe e0 she
+      | Just res <- splitIfold sht (simifold1 fe) (simifold1 e0) (simifold1 she)
+      -> res
+    e -> simrecurse simifold1 e
+
+simrecurse :: (forall env' a'. Exp env' a' -> Exp env' a') -> Exp env a -> Exp env a
+simrecurse f = \case
+    App a b -> App (f a) (f b)
+    Lam t e -> Lam t (f e)
+    Var t i -> Var t i
+    Let a e -> Let (f a) (f e)
+    Lit l -> Lit l
+    Cond a b c -> Cond (f a) (f b) (f c)
+    Const c -> Const c
+    Pair a b -> Pair (f a) (f b)
+    Fst e -> Fst (f e)
+    Snd e -> Snd (f e)
+    Build sht a b -> Build sht (f a) (f b)
+    Ifold sht a b c -> Ifold sht (f a) (f b) (f c)
+    Index a b -> Index (f a) (f b)
+    Shape e -> Shape (f e)
+    Undef t -> Undef t
+
+infixr :|
+data SimList = (forall env' a'. Exp env' a' -> Exp env' a') :| SimList
+             | SimEnd
+
+simfix :: SimList -> Exp env a -> Exp env a
+simfix list = \e -> let e' = looponce list e
+                    in case geq e e' of
+                         Just Refl -> e'
+                         Nothing -> simfix list e'
+  where
+    looponce :: SimList -> Exp env a -> Exp env a
+    looponce SimEnd e = e
+    looponce (f :| l) e = looponce l (f e)