-- The default type signature of type class methods are changed
-- to introduce a 'LiftDerived'* constraint and the same type class but on the 'Derived' repr,
-- this setup avoids to define the method with boilerplate code when its default
-- definition with 'liftDerived'* and 'derive' does what is expected by an instance
-- of the type class. This is almost as explained in:
-- https://ro-che.info/articles/2016-02-03-finally-tagless-boilerplate
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveGeneric #-} -- For NFData instances
{-# LANGUAGE DeriveAnyClass #-} -- For NFData instances
{-# LANGUAGE DeriveLift #-} -- For TH.Lift (Exception tok)
{-# LANGUAGE DerivingStrategies #-} -- For UnscopedRegister
{-# LANGUAGE StandaloneDeriving #-} -- For Show (Exception (InputToken inp))
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TemplateHaskell #-}
-- | Semantic of the grammar combinators used to express parsers,
-- in the convenient tagless-final encoding.
module Symantic.Parser.Grammar.Combinators where

import Data.Proxy (Proxy(..))
import Control.Monad (Monad(..))
import Control.DeepSeq (NFData(..))
import GHC.Generics (Generic)
-- import Data.Set (Set)
-- import GHC.TypeLits (KnownSymbol)
import Data.Bifunctor (second)
import Data.Bool (Bool(..), not, (||))
import Data.Char (Char)
import Data.Either (Either(..))
import Data.Eq (Eq(..))
import Data.Function ((.), flip, const, fix)
import Data.Ord (Ord(..), Ordering(..))
import Data.Int (Int)
import Data.Kind (Type, Constraint)
import Data.Maybe (Maybe(..))
import Data.String (String)
import Data.Semigroup (Semigroup(..))
import Text.Show (Show(..))
import Type.Reflection (Typeable, typeRep, eqTypeRep, (:~~:)(..), SomeTypeRep(..))
import qualified Data.Functor as Functor
import qualified Data.List as List
import qualified Language.Haskell.TH as TH
import qualified Language.Haskell.TH.Syntax as TH

import Symantic.Syntaxes.Derive
import qualified Symantic.Syntaxes.Classes as Prod
import Symantic.Parser.Grammar.Production

-- * Type 'ReprComb'
type ReprComb = Type -> Type

-- * Class 'CombAlternable'
class CombAlternable repr where
  -- | @('alt' es l r)@ parses @(l)@ and return its return value or,
  -- if it fails with an 'Exception' within @(es)@,
  -- parses @(r)@ from where @(l)@ has left the input stream,
  -- and returns its return value,
  -- otherwise throw the 'Exception' again.
  alt :: Exception -> repr a -> repr a -> repr a
  throw :: ExceptionLabel -> repr a
  -- | @('try' ra)@ records the input stream position,
  -- then parses like @(ra)@ and either returns its value it it succeeds or fails
  -- if it fails but with a reset of the input stream to the recorded position.
  -- Generally used on the first alternative: @('try' rl '<|>' rr)@.
  try :: repr a -> repr a
  default alt ::
    FromDerived2 CombAlternable repr =>
    Exception -> repr a -> repr a -> repr a
  default throw ::
    FromDerived CombAlternable repr =>
    ExceptionLabel -> repr a
  default try ::
    FromDerived1 CombAlternable repr =>
    repr a -> repr a
  alt = liftDerived2 . alt
  throw = liftDerived . throw
  try = liftDerived1 try

  -- | @(empty)@ parses nothing, always failing to return a value.
  empty :: repr a
  default empty ::
    FromDerived CombAlternable repr =>
    repr a
  empty = liftDerived empty

-- ** Type 'Exception'
data Exception
  =  ExceptionLabel ExceptionLabel
  |  ExceptionFailure
  -- |  ExceptionEnd
  deriving (Eq, Ord, Show, TH.Lift, Generic, NFData)
type ExceptionLabel = String
-- type Exceptions = Set Exception

-- | Like @('<|>')@ but with different returning types for the alternatives,
-- and a return value wrapped in an 'Either' accordingly.
(<+>) :: CombApplicable repr => CombAlternable repr => repr a -> repr b -> repr (Either a b)
p <+> q = Prod.left <$> p <|> Prod.right <$> q

(<|>) :: CombAlternable repr => repr a -> repr a -> repr a
(<|>) = alt ExceptionFailure

infixl 3 <|>, <+>

optionally :: CombApplicable repr => CombAlternable repr => repr a -> Production '[] b -> repr b
optionally p x = p $> x <|> pure x

optional :: CombApplicable repr => CombAlternable repr => repr a -> repr ()
optional = flip optionally (Prod.constant ())

option :: CombApplicable repr => CombAlternable repr => Production '[] a -> repr a -> repr a
option x p = p <|> pure x

choice :: CombAlternable repr => [repr a] -> repr a
choice = List.foldr (<|>) empty
 -- FIXME: Here hlint suggests to use Data.Foldable.asum,
 -- but at this point there is no asum for our own (<|>)

maybeP :: CombApplicable repr => CombAlternable repr => repr a -> repr (Maybe a)
maybeP p = option Prod.nothing (Prod.just <$> p)

manyTill :: CombApplicable repr => CombAlternable repr => repr a -> repr b -> repr [a]
manyTill p end = let go = end $> Prod.nil <|> p <:> go in go

-- * Class 'CombApplicable'
-- | This is like the usual 'Functor' and 'Applicative' type classes
-- from the @base@ package, but using @('Production' vs a)@ instead of just @(a)@
-- to be able to use and pattern match on some usual terms of type @(a)@ (like 'Prod.id')
-- and thus apply some optimizations.
-- @(repr)@, for "representation", is the usual tagless-final abstraction
-- over the many semantics that this syntax (formed by the methods
-- of type class like this one) will be interpreted.
class CombApplicable repr where
  -- | @(a2b '<$>' ra)@ parses like @(ra)@ but maps its returned value with @(a2b)@.
  (<$>) :: Production '[] (a -> b) -> repr a -> repr b
  (<$>) f = (pure f <*>)
  (<$>%) :: (Production '[] a -> Production '[] b) -> repr a -> repr b
  a2b <$>% ma = Prod.lam a2b <$> ma

  -- | Like '<$>' but with its arguments 'flip'-ped.
  (<&>) :: repr a -> Production '[] (a -> b) -> repr b
  (<&>) = flip (<$>)

  -- | @(a '<$' rb)@ parses like @(rb)@ but discards its returned value by replacing it with @(a)@.
  (<$) :: Production '[] a -> repr b -> repr a
  (<$) x = (pure x <*)

  -- | @(ra '$>' b)@ parses like @(ra)@ but discards its returned value by replacing it with @(b)@.
  ($>) :: repr a -> Production '[] b -> repr b
  ($>) = flip (<$)

  -- | @('pure' a)@ parses the empty string, always succeeding in returning @(a)@.
  pure :: Production '[] a -> repr a
  default pure ::
    FromDerived CombApplicable repr =>
    Production '[] a -> repr a
  pure = liftDerived . pure

  -- | @(ra2b '<*>' ra)@ parses sequentially @(ra2b)@ and then @(ra)@,
  -- and returns the application of the function returned by @(ra2b)@
  -- to the value returned by @(ra)@.
  (<*>) :: repr (a -> b) -> repr a -> repr b
  default (<*>) ::
    FromDerived2 CombApplicable repr =>
    repr (a -> b) -> repr a -> repr b
  (<*>) = liftDerived2 (<*>)

  -- | @(ra '<*' rb)@ parses sequentially @(ra)@ and then @(rb)@,
  -- and returns like @(ra)@, discarding the return value of @(rb)@.
  (<*) :: repr a -> repr b -> repr a
  (<*) = liftA2 Prod.const

  -- | @(ra '*>' rb)@ parses sequentially @(ra)@ and then @(rb)@,
  -- and returns like @(rb)@, discarding the return value of @(ra)@.
  (*>) :: repr a -> repr b -> repr b
  x *> y = (Prod.id <$ x) <*> y

  -- | Like '<*>' but with its arguments 'flip'-ped.
  (<**>) :: repr a -> repr (a -> b) -> repr b
  (<**>) = liftA2 (Prod.flip Prod..@ (Prod.$))
  {-
  (<**>) :: repr a -> repr (a -> b) -> repr b
  (<**>) = liftA2 (\a f -> f a)
  -}
  -- | @('liftA2' a2b2c ra rb)@ parses sequentially @(ra)@ and then @(rb)@,
  -- and returns the application of @(a2b2c)@ to the values returned by those parsers.
  liftA2 :: Production '[] (a -> b -> c) -> repr a -> repr b -> repr c
  liftA2 f x = (<*>) (f <$> x)

infixl 4 <*>, <*, *>, <**>


{-# INLINE (<:>) #-}
infixl 4 <:>
(<:>) :: CombApplicable repr => repr a -> repr [a] -> repr [a]
(<:>) = liftA2 Prod.cons

sequence :: CombApplicable repr => [repr a] -> repr [a]
sequence = List.foldr (<:>) (pure Prod.nil)

traverse :: CombApplicable repr => (a -> repr b) -> [a] -> repr [b]
traverse f = sequence . List.map f
 -- FIXME: Here hlint suggests to use Control.Monad.mapM,
 -- but at this point there is no mapM for our own sequence

repeat :: CombApplicable repr => Int -> repr a -> repr [a]
repeat n p = traverse (const p) [1..n]

between :: CombApplicable repr => repr o -> repr c -> repr a -> repr a
between open close p = open *> p <* close

void :: CombApplicable repr => repr a -> repr ()
void p = p *> unit

unit :: CombApplicable repr => repr ()
unit = pure (Prod.constant ())

-- * Class 'CombFoldable'
class CombFoldable repr where
  chainPre  :: repr (a -> a) -> repr a -> repr a
  chainPost :: repr a -> repr (a -> a) -> repr a
  chainPre  = liftDerived2 chainPre
  chainPost = liftDerived2 chainPost
  default chainPre ::
    FromDerived2 CombFoldable repr =>
    repr (a -> a) -> repr a -> repr a
  default chainPost ::
    FromDerived2 CombFoldable repr =>
    repr a -> repr (a -> a) -> repr a
  {-
  default chainPre ::
    CombApplicable repr =>
    CombAlternable repr =>
    repr (a -> a) -> repr a -> repr a
  default chainPost ::
    CombApplicable repr =>
    CombAlternable repr =>
    repr a -> repr (a -> a) -> repr a
  chainPre op p = go <*> p where go = (Prod..) <$> op <*> go <|> pure Prod.id
  chainPost p op = p <**> go where go = (Prod..) <$> op <*> go <|> pure Prod.id
  -}
  {-
  chainPre op p = flip (foldr ($)) <$> many op <*> p
  chainPost p op = foldl' (flip ($)) <$> p <*> many op
  -}

{-
conditional :: CombSelectable repr => [(Production '[] (a -> Bool), repr b)] -> repr a -> repr b -> repr b
conditional cs p def = match p fs qs def
  where (fs, qs) = List.unzip cs
-}

-- Parser Folds
pfoldr ::
 CombApplicable repr => CombFoldable repr =>
 Production '[] (a -> b -> b) -> Production '[] b -> repr a -> repr b
pfoldr f k p = chainPre (f <$> p) (pure k)

pfoldr1 ::
 CombApplicable repr => CombFoldable repr =>
 Production '[] (a -> b -> b) -> Production '[] b -> repr a -> repr b
pfoldr1 f k p = f <$> p <*> pfoldr f k p

pfoldl ::
 CombApplicable repr => CombFoldable repr =>
 Production '[] (b -> a -> b) -> Production '[] b -> repr a -> repr b
pfoldl f k p = chainPost (pure k) ((Prod.flip <$> pure f) <*> p)

pfoldl1 ::
 CombApplicable repr => CombFoldable repr =>
 Production '[] (b -> a -> b) -> Production '[] b -> repr a -> repr b
pfoldl1 f k p = chainPost (f <$> pure k <*> p) ((Prod.flip <$> pure f) <*> p)

-- Chain Combinators
chainl1' ::
 CombApplicable repr => CombFoldable repr =>
 Production '[] (a -> b) -> repr a -> repr (b -> a -> b) -> repr b
chainl1' f p op = chainPost (f <$> p) (Prod.flip <$> op <*> p)

chainl1 ::
 CombApplicable repr => CombFoldable repr =>
 repr a -> repr (a -> a -> a) -> repr a
chainl1 = chainl1' Prod.id

{-
chainr1' :: ParserOps rep => rep (a -> b) -> repr a -> repr (a -> b -> b) -> repr b
chainr1' f p op = newRegister_ Prod.id $ \acc ->
  let go = bind p $ \x ->
           modify acc (Prod.flip (Prod..@) <$> (op <*> x)) *> go
       <|> f <$> x
  in go <**> get acc

chainr1 :: repr a -> repr (a -> a -> a) -> repr a
chainr1 = chainr1' Prod.id

chainr :: repr a -> repr (a -> a -> a) -> Production '[] a -> repr a
chainr p op x = option x (chainr1 p op)
-}

chainl ::
 CombApplicable repr => CombAlternable repr => CombFoldable repr =>
 repr a -> repr (a -> a -> a) -> Production '[] a -> repr a
chainl p op x = option x (chainl1 p op)

-- Derived Combinators
many ::
 CombApplicable repr => CombFoldable repr =>
 repr a -> repr [a]
many = pfoldr Prod.cons Prod.nil

manyN ::
 CombApplicable repr => CombFoldable repr =>
 Int -> repr a -> repr [a]
manyN n p = List.foldr (const (p <:>)) (many p) [1..n]

some ::
 CombApplicable repr => CombFoldable repr =>
 repr a -> repr [a]
some = manyN 1

skipMany ::
 CombApplicable repr => CombFoldable repr =>
 repr a -> repr ()
--skipMany p = let skipManyp = p *> skipManyp <|> unit in skipManyp
skipMany = void . pfoldl Prod.const Prod.unit -- the void here will encourage the optimiser to recognise that the register is unused

skipManyN ::
 CombApplicable repr => CombFoldable repr =>
 Int -> repr a -> repr ()
skipManyN n p = List.foldr (const (p *>)) (skipMany p) [1..n]

skipSome ::
 CombApplicable repr => CombFoldable repr =>
 repr a -> repr ()
skipSome = skipManyN 1

sepBy ::
 CombApplicable repr => CombAlternable repr => CombFoldable repr =>
 repr a -> repr b -> repr [a]
sepBy p sep = option Prod.nil (sepBy1 p sep)

sepBy1 ::
 CombApplicable repr => CombAlternable repr => CombFoldable repr =>
 repr a -> repr b -> repr [a]
sepBy1 p sep = p <:> many (sep *> p)

endBy ::
 CombApplicable repr => CombAlternable repr => CombFoldable repr =>
 repr a -> repr b -> repr [a]
endBy p sep = many (p <* sep)

endBy1 ::
 CombApplicable repr => CombAlternable repr => CombFoldable repr =>
 repr a -> repr b -> repr [a]
endBy1 p sep = some (p <* sep)

sepEndBy ::
 CombApplicable repr => CombAlternable repr => CombFoldable repr =>
 repr a -> repr b -> repr [a]
sepEndBy p sep = option Prod.nil (sepEndBy1 p sep)

sepEndBy1 ::
 CombApplicable repr => CombAlternable repr => CombFoldable repr =>
 repr a -> repr b -> repr [a]
sepEndBy1 p sep =
  let seb1 = p <**> (sep *> (Prod.flip Prod..@ Prod.cons <$> option Prod.nil seb1)
                 <|> pure (Prod.flip Prod..@ Prod.cons Prod..@ Prod.nil))
  in seb1

{-
sepEndBy1 :: repr a -> repr b -> repr [a]
sepEndBy1 p sep = newRegister_ Prod.id $ \acc ->
  let go = modify acc ((Prod.flip (Prod..)) Prod..@ Prod.cons <$> p)
         *> (sep *> (go <|> get acc) <|> get acc)
  in go <*> pure Prod.nil
-}

-- * Class 'CombMatchable'
class CombMatchable repr where
  conditional ::
    repr a -> [(Production '[] (a -> Bool), repr b)] -> repr b -> repr b
  conditional a bs = liftDerived1
    (conditional (derive a) (second derive Functor.<$> bs))
  default conditional ::
    FromDerived1 CombMatchable repr => Derivable repr =>
    repr a -> [(Production '[] (a -> Bool), repr b)] -> repr b -> repr b

match ::
  CombMatchable repr =>
  Eq a => TH.Lift a =>
  repr a -> [Production '[] a] -> (Production '[] a -> repr b) -> repr b -> repr b
match a as p = conditional a
  ((\v ->
    ( Prod.lam (v Prod.==)
    , p v
    )
  ) Functor.<$> as)

predicate ::
  CombMatchable repr =>
  Production '[] (a -> Bool) -> repr a -> repr b -> repr b -> repr b
predicate p a b = conditional a [(p, b)]

infixl 4 <?:>
(<?:>) ::
  CombMatchable repr =>
  repr Bool -> (repr a, repr a) -> repr a
cond <?:> (p, q) = predicate Prod.id cond p q

-- * Class 'CombSatisfiable'
class CombSatisfiable tok repr where
  satisfy :: Production '[] (tok -> Bool) -> repr tok
  default satisfy ::
    FromDerived (CombSatisfiable tok) repr =>
    Production '[] (tok -> Bool) -> repr tok
  satisfy = liftDerived . satisfy

-- | Global 'TH.Name' to refer to the @(InputToken inp)@ type
-- from TemplateHaskell code.
inputTokenProxy :: TH.Name
inputTokenProxy = TH.mkName "inputToken"

char ::
  CombApplicable repr =>
  CombSatisfiable Char repr =>
  Char -> repr Char
char c = satisfy (Prod.equal Prod..@ Prod.constant c) $> Prod.constant c

item :: forall tok repr.
  Ord tok => Show tok => Typeable tok => TH.Lift tok => NFData tok =>
  CombSatisfiable tok repr => repr tok
item = satisfy (Prod.const Prod..@ Prod.constant True)

anyChar ::
  CombAlternable repr =>
  CombSatisfiable Char repr =>
  repr Char
anyChar = item

string ::
  CombApplicable repr => CombAlternable repr =>
  CombSatisfiable Char repr =>
  [Char] -> repr [Char]
string = try . traverse char

oneOf ::
  Ord tok => Show tok => Typeable tok => TH.Lift tok => NFData tok =>
  CombSatisfiable tok repr =>
  [tok] -> repr tok
oneOf ts = satisfy (production (`List.elem` ts) [||\t -> $$(ofChars ts [||t||])||])

noneOf ::
  (Ord tok, Show tok, TH.Lift tok, NFData tok, Typeable tok) =>
  CombSatisfiable tok repr =>
  [tok] -> repr tok
noneOf ts = satisfy (production (not . (`List.elem` ts)) [||\c -> not $$(ofChars ts [||c||])||])

ofChars ::
  TH.Lift tok => Eq tok =>
  {-alternatives-}[tok] ->
  {-input-}TH.CodeQ tok ->
  TH.CodeQ Bool
ofChars = List.foldr
  (\tok acc inp -> [|| tok == $$inp || $$(acc inp) ||])
  (const [||False||])

more ::
  CombAlternable repr =>
  CombApplicable repr =>
  CombSatisfiable Char repr =>
  CombLookable repr => repr ()
more = look (void (item @Char))

token ::
  (Ord tok, Show tok, TH.Lift tok, NFData tok, Typeable tok) =>
  CombAlternable repr =>
  CombApplicable repr =>
  CombSatisfiable tok repr =>
  tok -> repr tok
token tok = satisfy (Prod.equal Prod..@ Prod.constant tok) $> Prod.constant tok
-- token tok = satisfy [ExceptionToken tok] (Prod.eq Prod..@ Prod.qual Prod..@ Prod.char tok) $> Prod.char tok

tokens ::
  (Ord tok, Show tok, TH.Lift tok, NFData tok, Typeable tok) =>
  CombApplicable repr => CombAlternable repr =>
  CombSatisfiable tok repr => [tok] -> repr [tok]
tokens = try . traverse token

-- * Class 'CombSelectable'
class CombSelectable repr where
  branch :: repr (Either a b) -> repr (a -> c) -> repr (b -> c) -> repr c
  default branch ::
    FromDerived3 CombSelectable repr =>
    repr (Either a b) -> repr (a -> c) -> repr (b -> c) -> repr c
  branch = liftDerived3 branch

when ::
  CombMatchable repr =>
  CombSelectable repr =>
  Prod.Constantable () repr =>
  repr Bool -> repr () -> repr ()
when p q = p <?:> (q, Prod.constant ())

while ::
  CombMatchable repr =>
  CombSelectable repr =>
  Prod.Constantable () repr =>
  repr Bool -> repr ()
while x = fix (when x)

-- * Class 'CombLookable'
class CombLookable repr where
  look :: repr a -> repr a
  -- |
  -- Note: following [Error Reporting in Parsing Expression Grammars](https://arxiv.org/abs/1405.6646v3)
  -- what happens inside a 'negLook' does not take part in error reporting at
  -- all, which is the simplest approach, and also gives a consistent result
  -- for: @(negLook . negLook)@.
  negLook :: repr a -> repr ()
  default look ::
    FromDerived1 CombLookable repr =>
    repr a -> repr a
  default negLook ::
    FromDerived1 CombLookable repr =>
    repr a -> repr ()
  look = liftDerived1 look
  negLook = liftDerived1 negLook

  eof :: repr ()
  eof = liftDerived eof
  default eof ::
    FromDerived CombLookable repr =>
    repr ()
  -- eof = negLook (satisfy @Char (Prod.const Prod..@ Prod.bool True))
             -- (item @Char)

-- Composite Combinators
-- someTill :: repr a -> repr b -> repr [a]
-- someTill p end = negLook end *> (p <:> manyTill p end)

{-
constp :: CombApplicable repr => repr a -> repr (b -> a)
constp = (Prod.const <$>)


-- Alias Operations
infixl 1 >>
(>>) :: CombApplicable repr => repr a -> repr b -> repr b
(>>) = (*>)

-- Monoidal Operations

infixl 4 <~>
(<~>) :: CombApplicable repr => repr a -> repr b -> repr (a, b)
(<~>) = liftA2 (Prod.runtime (,))

infixl 4 <~
(<~) :: CombApplicable repr => repr a -> repr b -> repr a
(<~) = (<*)

infixl 4 ~>
(~>) :: CombApplicable repr => repr a -> repr b -> repr b
(~>) = (*>)

-- Lift Operations
liftA2 ::
 CombApplicable repr =>
 Production '[] (a -> b -> c) -> repr a -> repr b -> repr c
liftA2 f x = (<*>) (fmap f x)

liftA3 ::
 CombApplicable repr =>
 Production '[] (a -> b -> c -> d) -> repr a -> repr b -> repr c -> repr d
liftA3 f a b c = liftA2 f a b <*> c

-}

{-
-- Combinators interpreters for 'Sym.Any'.
instance CombApplicable repr => CombApplicable (Sym.Any repr)
instance CombSatisfiable repr => CombSatisfiable (Sym.Any repr)
instance CombAlternable repr => CombAlternable (Sym.Any repr)
instance CombSelectable repr => CombSelectable (Sym.Any repr)
instance CombMatchable repr => CombMatchable (Sym.Any repr)
instance CombLookable repr => CombLookable (Sym.Any repr)
instance CombFoldable repr => CombFoldable (Sym.Any repr)
-}

-- * Type 'Register'
newtype Register r a = Register { unRegister :: UnscopedRegister a }
  deriving (Eq, Show)

-- ** Type 'UnscopedRegister'
newtype UnscopedRegister r = UnscopedRegister { unUnscopedRegister :: TH.Name }
  deriving (Eq)
  deriving newtype Show


{-
put_ :: ParserOps rep => Register r a -> rep a -> Parser ()
put_ r = put r . pure

gets_ :: ParserOps rep => Register r a -> rep (a -> b) -> Parser b
gets_ r = gets r . pure

modify_ :: ParserOps rep => Register r a -> rep (a -> a) -> Parser ()
modify_ r = modify r . pure
-}

gets ::
  CombApplicable repr =>
  CombRegisterable repr =>
  Register r a -> repr (a -> b) -> repr b
gets r p = p <*> get r

modify ::
  CombApplicable repr =>
  CombRegisterable repr =>
  Register r a -> repr (a -> a) -> repr ()
modify r p = put r (gets r p)

move ::
  CombRegisterable repr =>
  Register r1 a -> Register r2 a -> repr ()
move dst src = put dst (get src)

bind ::
  CombRegisterable repr =>
  repr a -> (repr a -> repr b) -> repr b
bind p f = new p (f . get)

local ::
  CombApplicable repr =>
  CombRegisterable repr =>
  Register r a -> repr a -> repr b -> repr b
local r p q = bind (get r) (\x -> put r p *> q <* put r x)

swap ::
  CombApplicable repr =>
  CombRegisterable repr =>
  Register r1 a -> Register r2 a -> repr ()
swap r1 r2 = bind (get r1) (\x -> move r1 r2 *> put r2 x)

rollback ::
  CombAlternable repr =>
  CombApplicable repr =>
  CombRegisterable repr =>
  Register r a -> repr b -> repr b
rollback r p = bind (get r) (\x -> p <|> put r x *> empty)

for ::
  CombApplicable repr =>
  CombMatchable repr =>
  CombSelectable repr =>
  CombRegisterable repr =>
  Prod.Constantable () repr =>
  repr a -> repr (a -> Bool) -> repr (a -> a) -> repr () -> repr ()
for init cond step body =
  new init (\i ->
    let cond' = gets i cond in
    when cond' (while (body *> modify i step *> cond'))
  )


-- ** Class 'CombRegisterable'
class CombRegisterable (repr::ReprComb) where
  new :: repr a -> (forall r. Register r a -> repr b) -> repr b
  get :: Register r a -> repr a
  put :: Register r a -> repr a -> repr ()
  default new ::
    FromDerived CombRegisterable repr => Derivable repr =>
    repr a -> (forall r. Register r a -> repr b) -> repr b
  default get ::
    FromDerived CombRegisterable repr =>
    Register r a -> repr a
  default put ::
    FromDerived1 CombRegisterable repr =>
    Register r a -> repr a -> repr ()
  new ini f = liftDerived (new (derive ini) (derive . f))
  get = liftDerived . get
  put = liftDerived1 . put