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