module StateFunctorInstances where
import Tools(swap23, assocpairl, assocpairr, eitherout, eitherin)
import ArrowClasses
import ArrowTools(StateArrT(..), updateStateArrT)

instance Arrow q => Arrow (StateArrT s q) where
  arr f   = StateArrT (arr (first f))
  -- or lift (arr f), which is equal to StateArrT (first (arr f))
  StateArrT f >>> StateArrT g = StateArrT (f >>> g)
  StateArrT f <<< StateArrT g = StateArrT (f <<< g)
  first (StateArrT f) = f `seq`
    StateArrT (arr swap23 >>> first f >>> arr swap23)
  second (StateArrT f) = f `seq`
    StateArrT (arr assocpairl >>> second f >>> arr assocpairr)
    --       ((c,a),s) -> (c,(a,s)) -> (c,(b,s)) -> ((c,b),s)

instance ArrowChoice q => ArrowChoice (StateArrT s q) where
  StateArrT f ||| StateArrT g = StateArrT (arr eitherout >>> (f ||| g))

instance ArrowZero q => ArrowZero (StateArrT s q) where
  zeroA = StateArrT zeroA

instance ArrowPlus q => ArrowPlus (StateArrT s q) where
  StateArrT f <|> StateArrT g = StateArrT ((f <|> g) >>> arr eitherin)
  StateArrT f <+> StateArrT g = StateArrT (f <+> g)

instance Arrow q => StateArrow s (StateArrT s q) where
  update f = updateStateArrT (arr f)
  --updateAndReturn f = StateArrT (arr (\(~(a,s))->(s,f s)))

instance Arrow q => ArrowTransformer (StateArrT s) q where
  lift f = StateArrT (first f)