{-# LANGUAGE GADTs, KindSignatures, TypeOperators, PatternGuards , ScopedTypeVariables, TypeFamilies, FlexibleInstances , MultiParamTypeClasses, FlexibleContexts #-} {-# OPTIONS_GHC -Wall -fno-warn-missing-methods -fno-warn-missing-signatures #-} module Exp where import Control.Applicative (pure,(<$>),(<*>)) import qualified Data.IntMap as I import Data.Maybe (fromMaybe) -- ,fromJust -- package ty import Data.Proof.EQ -- Import one of the following but not both import Data.Ty -- import CustomTy import Data.TReify {-------------------------------------------------------------------- Expressions --------------------------------------------------------------------} data Op :: * -> * where Add :: Num a => Op (a -> a -> a) Mul :: Num a => Op (a -> a -> a) Lit :: Show a => a -> Op a instance ShowF Op where showF Add = "Add" showF Mul = "Mul" showF (Lit a) = show a instance Show (Op a) where show = showF -- Expressions, parameterized by variable type (constructor) and -- expression type. data E :: (* -> *) -> * -> * where Op :: Op a -> E v a (:^) :: (Typeable a, Typeable b) => E v (a -> b) -> E v a -> E v b Let :: v a -> E v a -> E v b -> E v b Var :: v a -> E v a data N :: (* -> *) -> * -> * where ON :: Op a -> N v a App :: (Typeable a, Typeable b) => v (a -> b) -> v a -> N v b instance ShowF v => ShowF (N v) where showF (ON o) = unwords ["ON" ,showF o] showF (App a b) = unwords ["App",showF a,showF b] instance Show (E (V Ty) a) where show (Op o) = show o show (u :^ v) = parens $ unwords [show u,show v] show (Let v a b) = unwords ["let",show v,"=",show a,"in",show b] show (Var v) = show v parens :: String -> String parens = ("(" ++) . (++ ")") instance MuRef Ty (E v) where type DeRef (E v) = N mapDeRef _ _ (Op o) = pure $ ON o mapDeRef k _ (f :^ a) = App <$> k ty f <*> k ty a mapDeRef _ _ Let{} = notSupp "Let" mapDeRef _ _ Var{} = notSupp "Var" notSupp :: String -> a notSupp meth = error $ "mapDeRef on E: "++meth++" not supported" -- TODO: Consider splitting Let/Var off from E. Then E wouldn't need the -- v parameter. -- -- I could use N and Mu to define an E without Let & Var and then use a -- standard MuRef instance and a standard extension to include Let & Var. -- Wouldn't be as convenient for simplification rules, because of the -- explicit Mu constructors. Consider it, however, since it makes a -- simple type distinction between the input to CSE and the output. Oh, -- and I could have different ways to embed Let. For C-like languages, I -- could restrict the lets to be on the outside (odd for conditionals, -- though). nodeE :: N v a -> E v a nodeE (ON o) = Op o nodeE (App u v) = Var u :^ Var v unGraph :: Typeable a => Graph Ty N a -> E (V Ty) a unGraph (Graph binds root) = foldr (\ (Bind v n) -> Let v (nodeE n)) (Var root) (reverse binds) -- Convert expressions to simple SSA forms ssa :: Typeable a => E (V Ty) a -> IO (E (V Ty) a) ssa = fmap unGraph . reifyGraph ty children :: N (V Ty) a -> [Id] children (ON _) = [] children (App (V a _) (V b _)) = [a,b] childrenB :: Bind Ty N -> [Id] childrenB (Bind _ n) = children n -- Number of references for each node. Important: partially apply, so -- that the binding list can be converted just once into an efficiently -- searchable representation. uses :: [Bind Ty N] -> (Id -> Int) uses = fmap (fromMaybe 0) . flip I.lookup . histogram . concatMap childrenB -- histogram :: Ord k => [k] -> I.Map k Int -- histogram = foldr (\ k -> I.insertWith (+) k 1) I.empty histogram :: [Int] -> I.IntMap Int histogram = foldr (\ k -> I.insertWith (+) k 1) I.empty -- bindsF :: [Bind v] -> (Ty a -> v a -> N v a) -- bindsF = undefined -- Slow version bindsF' :: [Bind Ty N] -> (V Ty a -> N (V Ty) a) bindsF' [] _ = error "bindsF: ran out of bindings" bindsF' (Bind (V i a) n : binds') v@(V i' a') | Just Refl <- a `tyEq` a', i == i' = n | otherwise = bindsF' binds' v -- Fast version, using an IntMap. Important: partially apply. bindsF :: forall n a. [Bind Ty n] -> (V Ty a -> n (V Ty) a) bindsF binds = \ (V i' a') -> extract a' (I.lookup i' m) where m :: I.IntMap (Bind Ty n) m = I.fromList [(i,b) | b@(Bind (V i _) _) <- binds] extract :: Ty a' -> Maybe (Bind Ty n) -> n (V Ty) a' extract _ Nothing = error "bindsF: variable not found" extract a' (Just (Bind (V _ a) n)) | Just Refl <- a `tyEq` a' = n | otherwise = error "bindsF: wrong type" unGraph2 :: Graph Ty N a -> E (V Ty) a unGraph2 (Graph binds root) = foldr llet (var' root) (reverse binds) where -- Wrap a let if non-trivial llet :: Bind Ty N -> E (V Ty) b -> E (V Ty) b llet bind | trivial bind = id llet (Bind v n) = Let v (nodeE' n) -- How many uses of variable count :: Id -> Int count = uses binds -- Bindings as IntMap lookup psf :: V Ty a -> N (V Ty) a psf = bindsF binds -- Too trivial to bother abstracting. trivial :: Bind Ty N -> Bool trivial (Bind _ (ON _)) = True trivial (Bind (V i _) _) | count i < 2 = True trivial _ = False -- Like nodeE but with inlining of trivial bindings nodeE' :: N (V Ty) a -> E (V Ty) a nodeE' (ON o) = Op o nodeE' (App a b) = var' a :^ var' b -- Variable reference or inline var' :: V Ty a -> E (V Ty) a var' v | trivial (Bind v n) = nodeE' n | otherwise = Var v where n = psf v -- TODO: generalize unGraph2 from V. -- | Common subexpression elimination cse :: Typeable a => E (V Ty) a -> IO (E (V Ty) a) cse = fmap unGraph2 . reifyGraph ty {-------------------------------------------------------------------- Utilities for convenient expression building --------------------------------------------------------------------} op2 :: (Typeable a, Typeable b, Typeable c) => Op (a -> b -> c) -> E v a -> E v b -> E v c op2 h a b = Op h :^ a :^ b instance Eq (E v a) instance (Typeable a, Num a) => Num (E (V Ty) a) where fromInteger x = Op (Lit (fromInteger x)) (+) = op2 Add (*) = op2 Mul sqr :: Num a => a -> a sqr x = x * x {-------------------------------------------------------------------- Testing --------------------------------------------------------------------} -- type-specialize reify :: (MuRef Ty h, Typeable a) => h a -> IO (Graph Ty (DeRef h) a) reify = reifyGraph ty -- test expressions e1 = 3 + 5 :: E (V Ty) Integer e2 = e1 * e1 e3 = 3 + 3 :: E (V Ty) Integer {- > e1 ((Add 3) 5) > reify e1 let [x0 = App x1 x4,x4 = ON 5,x1 = App x2 x3,x3 = ON 3,x2 = ON Add] in x0 > ssa e1 let x2 = Add in let x3 = 3 in let x1 = (x2 x3) in let x4 = 5 in let x0 = (x1 x4) in x0 > cse e1 ((Add 3) 5) > e2 ((Mul ((Add 3) 5)) ((Add 3) 5)) > reify e2 let [x0 = App x1 x3,x1 = App x2 x3,x3 = App x4 x7,x7 = ON 5,x4 = App x5 x6,x6 = ON 3,x5 = ON Add,x2 = ON Mul] in x0 > ssa e2 let x2 = Mul in let x5 = Add in let x6 = 3 in let x4 = (x5 x6) in let x7 = 5 in let x3 = (x4 x7) in let x1 = (x2 x3) in let x0 = (x1 x3) in x0 > cse e2 let x3 = ((Add 3) 5) in ((Mul x3) x3) > e3 ((Add 3) 3) > reify e3 let [x0 = App x1 x3,x1 = App x2 x3,x3 = ON 3,x2 = ON Add] in x0 > > ssa e3 let x2 = Add in let x3 = 3 in let x1 = (x2 x3) in let x0 = (x1 x3) in x0 > cse e3 ((Add 3) 3) -} test :: Int -> E (V Ty) Integer test n = iterate sqr e1 !! n {- > test 2 ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5))) > reify (test 2) let [x0 = App x1 x3,x1 = App x2 x3,x3 = App x4 x6,x4 = App x5 x6,x6 = App x7 x10,x10 = ON 5,x7 = App x8 x9,x9 = ON 3,x8 = ON Add,x5 = ON Mul,x2 = ON Mul] in x0 > ssa (test 2) let x2 = Mul in let x5 = Mul in let x8 = Add in let x9 = 3 in let x7 = (x8 x9) in let x10 = 5 in let x6 = (x7 x10) in let x4 = (x5 x6) in let x3 = (x4 x6) in let x1 = (x2 x3) in let x0 = (x1 x3) in x0 > cse (test 2) let x6 = ((Add 3) 5) in let x3 = ((Mul x6) x6) in ((Mul x3) x3) > test 5 ((Mul ((Mul ((Mul ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))) ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5))))) ((Mul ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))) ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))))) ((Mul ((Mul ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))) ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5))))) ((Mul ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))) ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))))) > reify (test 5) let [x0 = App x1 x3,x1 = App x2 x3,x3 = App x4 x6,x4 = App x5 x6,x6 = App x7 x9,x7 = App x8 x9,x9 = App x10 x12,x10 = App x11 x12,x12 = App x13 x15,x13 = App x14 x15,x15 = App x16 x19,x19 = ON 5,x16 = App x17 x18,x18 = ON 3,x17 = ON Add,x14 = ON Mul,x11 = ON Mul,x8 = ON Mul,x5 = ON Mul,x2 = ON Mul] in x0 > ssa (test 5) let x2 = Mul in let x5 = Mul in let x8 = Mul in let x11 = Mul in let x14 = Mul in let x17 = Add in let x18 = 3 in let x16 = (x17 x18) in let x19 = 5 in let x15 = (x16 x19) in let x13 = (x14 x15) in let x12 = (x13 x15) in let x10 = (x11 x12) in let x9 = (x10 x12) in let x7 = (x8 x9) in let x6 = (x7 x9) in let x4 = (x5 x6) in let x3 = (x4 x6) in let x1 = (x2 x3) in let x0 = (x1 x3) in x0 > cse (test 5) let x15 = ((Add 3) 5) in let x12 = ((Mul x15) x15) in let x9 = ((Mul x12) x12) in let x6 = ((Mul x9) x9) in let x3 = ((Mul x6) x6) in ((Mul x3) x3) -} e4 = e1 ^ (29 :: Integer) {- > e4 ((Mul ((Mul ((Mul ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))) ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5))))) ((Mul ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))) ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))))) ((Mul ((Mul ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))) ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5))))) ((Mul ((Mul ((Mul ((Add 3) 5)) ((Add 3) 5))) ((Mul ((Add 3) 5)) ((Add 3) 5)))) ((Add 3) 5)))) > reify e4 let [x0 = App x1 x20,x20 = App x21 x23,x23 = App x24 x15,x24 = App x25 x9,x25 = ON Mul,x21 = App x22 x6,x22 = ON Mul,x1 = App x2 x3,x3 = App x4 x6,x4 = App x5 x6,x6 = App x7 x9,x7 = App x8 x9,x9 = App x10 x12,x10 = App x11 x12,x12 = App x13 x15,x13 = App x14 x15,x15 = App x16 x19,x19 = ON 5,x16 = App x17 x18,x18 = ON 3,x17 = ON Add,x14 = ON Mul,x11 = ON Mul,x8 = ON Mul,x5 = ON Mul,x2 = ON Mul] in x0 > ssa e4 let x2 = Mul in let x5 = Mul in let x8 = Mul in let x11 = Mul in let x14 = Mul in let x17 = Add in let x18 = 3 in let x16 = (x17 x18) in let x19 = 5 in let x15 = (x16 x19) in let x13 = (x14 x15) in let x12 = (x13 x15) in let x10 = (x11 x12) in let x9 = (x10 x12) in let x7 = (x8 x9) in let x6 = (x7 x9) in let x4 = (x5 x6) in let x3 = (x4 x6) in let x1 = (x2 x3) in let x22 = Mul in let x21 = (x22 x6) in let x25 = Mul in let x24 = (x25 x9) in let x23 = (x24 x15) in let x20 = (x21 x23) in let x0 = (x1 x20) in x0 > cse e4 let x15 = ((Add 3) 5) in let x12 = ((Mul x15) x15) in let x9 = ((Mul x12) x12) in let x6 = ((Mul x9) x9) in ((Mul ((Mul x6) x6)) ((Mul x6) ((Mul x9) x15))) -}