add "sanity" target to Makefile
[coq-hetmet.git] / examples / GArrowPortShape.hs
1 {-# LANGUAGE MultiParamTypeClasses, GADTs, FlexibleContexts, FlexibleInstances, TypeFamilies #-}
2 -----------------------------------------------------------------------------
3 -- |
4 -- Module      :  GArrowPortShape
5 -- Copyright   :  none
6 -- License     :  public domain
7 --
8 -- Maintainer  :  Adam Megacz <megacz@acm.org>
9 -- Stability   :  experimental
10 --
11 -- | We cannot, at run time, query to find out the input and output
12 -- port types of a GArrowSkeleton since Haskell erases types during
13 -- compilation.  Using Data.Typeable is problematic here because
14 -- GAS_comp and GAS_loop{l,r} have an existential type.
15 --
16 -- In spite of this, we can determine the "shape" of the ports --
17 -- which ports are of unit type, and which ports must be tensors.  A
18 -- GArrowPortShape is a GArrowSkeleton along with this
19 -- information for certain nodes (the inference mechanism below adds
20 -- it on every node).
21 --
22 module GArrowPortShape (GArrowPortShape(..), PortShape(..), detectShape)
23 where
24 import Prelude hiding ( id, (.), lookup )
25 import Control.Category
26 import GHC.HetMet.GArrow
27 import Unify
28 import GArrowSkeleton
29 import Control.Monad.State
30
31 --
32 -- | Please keep in mind that the "shapes" computed below are simply the
33 -- least-complicated shapes that could possibly work.  Just because a
34 -- GArrowPortShape has an input port of shape (x,y)
35 -- doesn't mean it couldn't later be used in a context where its input
36 -- port had shape ((a,b),y)!  However, you can be assured that it
37 -- won't be used in a context where the input port has shape ().
38 --
39 data PortShape a = PortUnit
40                  | PortTensor (PortShape a) (PortShape a)
41                  | PortFree a
42
43 data GArrowPortShape m s a b =
44     GASPortPassthrough
45       (PortShape s)
46       (PortShape s)
47       (m a b)
48   | GASPortShapeWrapper
49       (PortShape s)
50       (PortShape s)
51       (GArrowSkeleton (GArrowPortShape m s) a b)
52
53 --
54 -- implementation below; none of this is exported
55 --
56
57 type UPort = PortShape UVar
58
59 instance Unifiable UPort where
60   unify' (PortTensor x1 y1) (PortTensor x2 y2) = mergeU (unify x1 x2) (unify y1 y2)
61   unify' PortUnit PortUnit                     = emptyUnifier
62   unify' s1 s2                                 = error $ "Unifiable UPort got impossible unification case: "
63 --                                                          ++ show s1 ++ " and " ++ show s2
64   inject                                       = PortFree
65   project (PortFree v)                         = Just v
66   project _                                    = Nothing
67   occurrences (PortFree v)                     = [v]
68   occurrences (PortTensor x y)                 = occurrences x ++ occurrences y
69   occurrences PortUnit                         = []
70
71 -- detection monad
72 type DetectM a = State ((Unifier UPort),[UVar]) a
73
74 shapes :: GArrowPortShape m UVar a b -> (UPort,UPort)
75 shapes (GASPortPassthrough  x y _) = (x,y)
76 shapes (GASPortShapeWrapper x y _) = (x,y)
77
78 unifyM :: UPort -> UPort -> DetectM ()
79 unifyM p1 p2 = do { (u,vars) <- get
80                   ; put (mergeU u $ unify p1 p2 , vars)
81                   }
82
83 freshM :: DetectM UVar
84 freshM = do { (u,(v:vars)) <- get
85             ; put (u,vars)
86             ; return v
87             }
88
89 -- recursive version of getU
90 getU' :: Unifier UPort -> UPort -> PortShape ()
91 getU' u (PortTensor x y)  = PortTensor (getU' u x) (getU' u y)
92 getU' _ PortUnit          = PortUnit
93 getU' u x@(PortFree v)    = case Unify.getU u v  of
94                                      Nothing -> PortFree () -- or x
95                                      Just x' -> getU' u x'
96
97 resolveG :: Unifier UPort -> (GArrowPortShape m UVar a b) -> GArrowPortShape m () a b
98 resolveG u (GASPortPassthrough  x y m) = GASPortPassthrough  (getU' u x) (getU' u y) m
99 resolveG u (GASPortShapeWrapper x y g) = GASPortShapeWrapper (getU' u x) (getU' u y) (resolveG' g)
100  where
101   resolveG' :: GArrowSkeleton (GArrowPortShape m UVar)             a b -> 
102                GArrowSkeleton (GArrowPortShape m ())   a b
103   resolveG' (GAS_id           ) = GAS_id
104   resolveG' (GAS_comp      f g) = GAS_comp (resolveG' f) (resolveG' g)
105   resolveG' (GAS_first       f) = GAS_first (resolveG' f)
106   resolveG' (GAS_second      f) = GAS_second (resolveG' f)
107   resolveG' GAS_cancell         = GAS_cancell
108   resolveG' GAS_cancelr         = GAS_cancelr
109   resolveG' GAS_uncancell       = GAS_uncancell
110   resolveG' GAS_uncancelr       = GAS_uncancelr
111   resolveG' GAS_drop            = GAS_drop
112   resolveG' (GAS_const i)       = GAS_const i
113   resolveG' GAS_copy            = GAS_copy
114   resolveG' GAS_merge           = GAS_merge
115   resolveG' GAS_swap            = GAS_swap
116   resolveG' GAS_assoc           = GAS_assoc
117   resolveG' GAS_unassoc         = GAS_unassoc
118   resolveG' (GAS_loopl f)       = GAS_loopl (resolveG' f)
119   resolveG' (GAS_loopr f)       = GAS_loopr (resolveG' f)
120   resolveG' (GAS_misc g )       = GAS_misc $ resolveG u g
121
122 detectShape :: GArrowSkeleton m a b -> GArrowPortShape m () a b
123 detectShape g = runM (detect g)
124
125 runM :: DetectM (GArrowPortShape m UVar a b) -> GArrowPortShape m () a b
126 runM f = let s     = (emptyUnifier,uvarSupply)
127              g     = evalState f s
128              (u,_) = execState f s
129           in resolveG u g
130
131 detect :: GArrowSkeleton m a b -> DetectM (GArrowPortShape m UVar a b)
132 detect (GAS_id      ) = do { x <- freshM ; return $ GASPortShapeWrapper (PortFree x) (PortFree x) GAS_id }
133 detect (GAS_comp f g) = do { f' <- detect f
134                            ; g' <- detect g
135                            ; unifyM (snd $ shapes f') (fst $ shapes g')
136                            ; return $ GASPortShapeWrapper (fst $ shapes f') (snd $ shapes g') (GAS_comp (GAS_misc f') (GAS_misc g'))
137                            }
138 detect (GAS_first  f) = do { x <- freshM
139                            ; f' <- detect f
140                            ; return $ GASPortShapeWrapper
141                                         (PortTensor (fst $ shapes f') (PortFree x))
142                                         (PortTensor (snd $ shapes f') (PortFree x))
143                                         (GAS_first (GAS_misc f'))
144                            }
145 detect (GAS_second f) = do { x <- freshM
146                            ; f' <- detect f
147                            ; return $ GASPortShapeWrapper
148                                         (PortTensor (PortFree x) (fst $ shapes f'))
149                                         (PortTensor (PortFree x) (snd $ shapes f'))
150                                         (GAS_second (GAS_misc f'))
151                            }
152 detect GAS_cancell    = do { x <- freshM; return$GASPortShapeWrapper (PortTensor PortUnit (PortFree x)) (PortFree x) GAS_cancell }
153 detect GAS_cancelr    = do { x <- freshM; return$GASPortShapeWrapper (PortTensor (PortFree x) PortUnit) (PortFree x) GAS_cancelr }
154 detect GAS_uncancell  = do { x <- freshM; return$GASPortShapeWrapper (PortFree x) (PortTensor PortUnit (PortFree x)) GAS_uncancell }
155 detect GAS_uncancelr  = do { x <- freshM; return$GASPortShapeWrapper (PortFree x) (PortTensor (PortFree x) PortUnit) GAS_uncancelr }
156 detect GAS_drop       = do { x <- freshM; return$GASPortShapeWrapper (PortFree x) PortUnit GAS_drop }
157 detect GAS_copy       = do { x <- freshM
158                            ; return $ GASPortShapeWrapper (PortFree x) (PortTensor (PortFree x) (PortFree x)) GAS_copy }
159 detect GAS_swap       = do { x <- freshM
160                            ; y <- freshM
161                            ; let x' = PortFree x
162                            ; let y' = PortFree y
163                            ; return $ GASPortShapeWrapper (PortTensor x' y') (PortTensor y' x') GAS_swap
164                            }
165 detect GAS_assoc      = do { x <- freshM; y <- freshM; z <- freshM
166                            ; let x' = PortFree x
167                            ; let y' = PortFree y
168                            ; let z' = PortFree z
169                            ; return $ GASPortShapeWrapper
170                                         (PortTensor (PortTensor x' y') z')
171                                         (PortTensor x' (PortTensor y' z'))
172                                         GAS_assoc
173                            }
174 detect GAS_unassoc    = do { x <- freshM; y <- freshM; z <- freshM
175                            ; let x' = PortFree x
176                            ; let y' = PortFree y
177                            ; let z' = PortFree z
178                            ; return $ GASPortShapeWrapper
179                                         (PortTensor x' (PortTensor y' z'))
180                                         (PortTensor (PortTensor x' y') z')
181                                         GAS_unassoc
182                            }
183 detect (GAS_const i)  = do { x <- freshM; return $ GASPortShapeWrapper PortUnit (PortFree x) (GAS_const i) }
184
185 -- FIXME: I need to fix the occurs check before I can make these different again
186 detect GAS_merge      = do { x <- freshM
187                            ; y <- freshM
188                            ; return $ GASPortShapeWrapper (PortTensor (PortFree x) (PortFree y)) (PortFree x) GAS_merge }
189 detect (GAS_loopl f)  = error "not implemented"
190 detect (GAS_loopr f)  = error "not implemented"
191
192 detect (GAS_misc f)   = error "not implemented"
193