more efficient encoding of function types
[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(..))
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' _ _                                   = error "impossible"
62   inject                                       = PortFree
63   project (PortFree v)                         = Just v
64   project _                                    = Nothing
65   occurrences (PortFree v)                     = [v]
66   occurrences (PortTensor x y)                 = occurrences x ++ occurrences y
67   occurrences PortUnit                         = []
68
69 -- detection monad
70 type DetectM a = State ((Unifier UPort),[UVar]) a
71
72 shapes :: GArrowPortShape m UVar a b -> (UPort,UPort)
73 shapes (GASPortPassthrough  x y _) = (x,y)
74 shapes (GASPortShapeWrapper x y _) = (x,y)
75
76 unifyM :: UPort -> UPort -> DetectM ()
77 unifyM p1 p2 = do { (u,vars) <- get
78                   ; put (mergeU u $ unify p1 p2 , vars)
79                   }
80
81 freshM :: DetectM UVar
82 freshM = do { (u,(v:vars)) <- get
83             ; put (u,vars)
84             ; return v
85             }
86
87 -- recursive version of getU
88 getU' :: Unifier UPort -> UPort -> PortShape ()
89 getU' u (PortTensor x y)  = PortTensor (getU' u x) (getU' u y)
90 getU' _ PortUnit          = PortUnit
91 getU' u x@(PortFree v)    = case Unify.getU u v  of
92                                      Nothing -> PortFree () -- or x
93                                      Just x' -> getU' u x'
94
95 resolveG :: Unifier UPort -> (GArrowPortShape m UVar a b) -> GArrowPortShape m () a b
96 resolveG u (GASPortPassthrough  x y m) = GASPortPassthrough  (getU' u x) (getU' u y) m
97 resolveG u (GASPortShapeWrapper x y g) = GASPortShapeWrapper (getU' u x) (getU' u y) (resolveG' g)
98  where
99   resolveG' :: GArrowSkeleton (GArrowPortShape m UVar)             a b -> 
100                GArrowSkeleton (GArrowPortShape m ())   a b
101   resolveG' (GAS_id           ) = GAS_id
102   resolveG' (GAS_comp      f g) = GAS_comp (resolveG' f) (resolveG' g)
103   resolveG' (GAS_first       f) = GAS_first (resolveG' f)
104   resolveG' (GAS_second      f) = GAS_second (resolveG' f)
105   resolveG' GAS_cancell         = GAS_cancell
106   resolveG' GAS_cancelr         = GAS_cancelr
107   resolveG' GAS_uncancell       = GAS_uncancell
108   resolveG' GAS_uncancelr       = GAS_uncancelr
109   resolveG' GAS_drop            = GAS_drop
110   resolveG' (GAS_const i)       = GAS_const i
111   resolveG' GAS_copy            = GAS_copy
112   resolveG' GAS_merge           = GAS_merge
113   resolveG' GAS_swap            = GAS_swap
114   resolveG' GAS_assoc           = GAS_assoc
115   resolveG' GAS_unassoc         = GAS_unassoc
116   resolveG' (GAS_loopl f)       = GAS_loopl (resolveG' f)
117   resolveG' (GAS_loopr f)       = GAS_loopr (resolveG' f)
118   resolveG' (GAS_misc g )       = GAS_misc $ resolveG u g
119
120 runM :: DetectM (GArrowPortShape m UVar a b) -> GArrowPortShape m () a b
121 runM f = let s     = (emptyUnifier,uvarSupply)
122              g     = evalState f s
123              (u,_) = execState f s
124           in resolveG u g
125
126 detect :: GArrowSkeleton m a b -> DetectM (GArrowPortShape m UVar a b)
127 detect (GAS_id      ) = do { x <- freshM ; return $ GASPortShapeWrapper (PortFree x) (PortFree x) GAS_id }
128 detect (GAS_comp g f) = do { f' <- detect f
129                            ; g' <- detect g
130                            ; unifyM (snd $ shapes f') (fst $ shapes g')
131                            ; return $ GASPortShapeWrapper (fst $ shapes f') (snd $ shapes g') (GAS_comp (GAS_misc g') (GAS_misc f'))
132                            }
133 detect (GAS_first  f) = do { x <- freshM
134                            ; f' <- detect f
135                            ; return $ GASPortShapeWrapper
136                                         (PortTensor (fst $ shapes f') (PortFree x))
137                                         (PortTensor (snd $ shapes f') (PortFree x))
138                                         (GAS_first (GAS_misc f'))
139                            }
140 detect (GAS_second f) = do { x <- freshM
141                            ; f' <- detect f
142                            ; return $ GASPortShapeWrapper
143                                         (PortTensor (PortFree x) (fst $ shapes f'))
144                                         (PortTensor (PortFree x) (snd $ shapes f'))
145                                         (GAS_second (GAS_misc f'))
146                            }
147 detect GAS_cancell    = do { x <- freshM; return$GASPortShapeWrapper (PortTensor PortUnit (PortFree x)) (PortFree x) GAS_cancell }
148 detect GAS_cancelr    = do { x <- freshM; return$GASPortShapeWrapper (PortTensor (PortFree x) PortUnit) (PortFree x) GAS_cancelr }
149 detect GAS_uncancell  = do { x <- freshM; return$GASPortShapeWrapper (PortFree x) (PortTensor PortUnit (PortFree x)) GAS_uncancell }
150 detect GAS_uncancelr  = do { x <- freshM; return$GASPortShapeWrapper (PortFree x) (PortTensor (PortFree x) PortUnit) GAS_uncancelr }
151 detect GAS_drop       = do { x <- freshM; return$GASPortShapeWrapper (PortFree x) PortUnit GAS_drop }
152 detect GAS_copy       = do { x <- freshM
153                            ; return $ GASPortShapeWrapper (PortFree x) (PortTensor (PortFree x) (PortFree x)) GAS_copy }
154 detect GAS_swap       = do { x <- freshM
155                            ; y <- freshM
156                            ; let x' = PortFree x
157                            ; let y' = PortFree y
158                            ; return $ GASPortShapeWrapper (PortTensor x' y') (PortTensor y' x') GAS_swap
159                            }
160 detect GAS_assoc      = do { x <- freshM; y <- freshM; z <- freshM
161                            ; let x' = PortFree x
162                            ; let y' = PortFree y
163                            ; let z' = PortFree z
164                            ; return $ GASPortShapeWrapper
165                                         (PortTensor (PortTensor x' y') z')
166                                         (PortTensor x' (PortTensor y' z'))
167                                         GAS_assoc
168                            }
169 detect GAS_unassoc    = do { x <- freshM; y <- freshM; z <- freshM
170                            ; let x' = PortFree x
171                            ; let y' = PortFree y
172                            ; let z' = PortFree z
173                            ; return $ GASPortShapeWrapper
174                                         (PortTensor x' (PortTensor y' z'))
175                                         (PortTensor (PortTensor x' y') z')
176                                         GAS_unassoc
177                            }
178 detect (GAS_const i)  = do { x <- freshM; return $ GASPortShapeWrapper PortUnit (PortFree x) (GAS_const i) }
179 detect GAS_merge      = do { x <- freshM
180                            ; return $ GASPortShapeWrapper (PortTensor (PortFree x) (PortFree x)) (PortFree x) GAS_merge }
181 detect (GAS_loopl f)  = error "not implemented"
182 detect (GAS_loopr f)  = error "not implemented"
183
184 detect (GAS_misc f)   = error "not implemented"
185