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