implementation module ASDS.Lens

import StdEnv
import Data.Func
import Data.Functor
import Data.List
import Data.Maybe
import Control.Monad
from Control.Monad.State import instance Monad (StateT st m), instance Functor (StateT st m), instance pure (StateT st m), instance <*> (StateT st m), instance MonadTrans (StateT st)
import Control.Monad.Trans
import Control.Monad.Fail
from Data.Map import :: Map(..), get, put

import ASDS, ASDS.Source

lens :: (p -> p1) (LensRead m p r r1) (LensWrite m p w r1 w1) (sds m p1 r1 w1) -> Lens sds m p r w | TC r1 & TC w1 & TC p1
lens param mapr mapw lens = Lens {param=param,mapr=mapr,mapw=mapw,lens=lens}

instance read (Lens sds) | read sds
where
	read (Lens t=:{mapr=LensReadConst f}) p = Read <$> liftT (f p)
	read (Lens t=:{param,mapr=LensRead mapr,lens}) p = read lens (param p) >>= \v->case v of
		Reading s = pure $ Reading $ Lens {t & lens=s}
		Read r = Read <$> liftT (mapr p r)

instance write (Lens sds) | read sds & write sds
where
	//First read the child when necessary
	write (Lens t=:{param,mapw=LensWrite mapw,lens}) p w = read lens (param p) >>= \v->case v of
		Reading s = pure $ Writing $ Lens {t & lens=rwpair s lens}
		Read r = write (Lens {t & mapw=LensWriteConst \p w->mapw p w r}) p w
	//Then do the actual writing
	write (Lens t=:{param,mapw=LensWriteConst mapw,lens}) p w = liftT (mapw p w) >>= \v->case v of
		?None = pure $ Written ()
		?Just w = write lens (param p) w >>= \v->case v of
			Writing s = pure $ Writing $ Lens {t & lens=rwpair lens s}
			Written () = pure $ Written ()

mapRead :: (r -> m r`) (sds m p r w) -> Lens sds m p r` w | TC r` & TC r & TC w & TC p & Monad m
mapRead f sds = lens id (LensRead \p->f) (LensWrite \p w r->pure $ ?Just w) sds

mapWrite :: (w` r -> m (? w)) (sds m p r w) -> Lens sds m p r w` | TC r & TC w & TC w` & TC p & Monad m
mapWrite f sds = lens id (LensRead \p r->pure r) (LensWrite \_ w r->f w r) sds

translate :: (p -> ps) (sds m ps r w) -> Lens sds m p r w | TC r & TC w & TC p & TC ps & Monad m
translate param sds = lens param (LensRead \p r->pure r) (LensWriteConst \p w->pure $ ?Just w) sds

fromDynStore :: (sds m p Dynamic Dynamic) -> Lens sds m p r w | TC r & TC w & TC p & MonadFail m
fromDynStore sds = lens id (LensRead read) (LensWriteConst write) sds
where
	read p (r :: r^) = pure r
	read p r = fail $ "fromDynStore: " +++ typeMismatchError r (dynamic undef :: r^)
	write p w = pure $ ?Just (dynamic w)

toDynStore :: (sds m p r w) -> Lens sds m p Dynamic Dynamic | TC r & TC w & TC p & MonadFail m
toDynStore sds = lens id (LensRead read) (LensWriteConst write) sds
where
	read p r = pure (dynamic r)
	write p (w :: w^) = pure $ ?Just w
	write p w = fail $ "toDynStore: " +++ typeMismatchError w (dynamic undef :: w^)

typeMismatchError :: Dynamic Dynamic -> String
typeMismatchError a b = "type mismatch " +++ toString (typeCodeOfDynamic a) +++ " and" +++ toString (typeCodeOfDynamic b)

indexedStore :: (sds m p [a] [a]) -> Lens sds m (p, Int) a a | TC a & TC p & MonadFail m
indexedStore sds = lens param (LensRead read) (LensWrite write) sds
where
	param (p, _) = p
	read (_, idx) r = case r !? idx of
		?None = fail "indexedStore: index out of range"
		?Just a = pure a
	write (_, idx) w r = pure $ ?Just $ updateAt idx w r

assocStore :: (sds m p [(k, v)] [(k, v)]) -> Lens sds m (p, k) v v | TC v & TC k & Eq k & TC p & MonadFail m
assocStore sds = lens param (LensRead read) (LensWrite write) sds
where
	param (p, _) = p
	read (_, k) r = case lookup k r of
		?None = fail "assocStore: key not present"
		?Just a = pure a
	write (_, k) w r = pure $ ?Just [(k, w):filter ((<>) k o fst) r]

keyedStore :: (sds m p (Map k v) (Map k v)) -> Lens sds m (p, k) v v | TC v & TC k & < k & TC p & MonadFail m
keyedStore sds = lens param (LensRead read) (LensWrite write) sds
where
	param (p, _) = p
	read (_, k) r = case get k r of
		?None = fail "keyedStore: key not present"
		?Just a = pure a
	write (_, k) w r = pure $ ?Just $ put k w r