||| Resources is framework for creating Resource Dependent EDSLs as described in the 2020 ECOOP paper:
||| *A Framework for Resource Dependent EDSLs in a Dependently Typed Language*.
|||
||| > Idris’ Effects library demonstrates how to embed resource dependent algebraic effect handlers into a
||| > dependently-typed host language, providing run-time and compile-time based reasoning on type-level resources.
||| > Building upon this work, Resources is a framework for realising Embedded Domain Specific Languages (EDSLs) with
||| > type-systems that contain domain specific substructural properties. Differing from Effects, Resources allows a
||| > language’s substructural properties to be encoded in a resource that is associated with language variables. Such
||| > an association allows for multiple effect instances to be reasoned about autonomically and without explicit
||| > type-level declaration. Type-level predicates are used as proof that the language’s substructural properties
||| > hold. Several exemplar EDSLs are presented that illustrates our framework’s operation and how dependent types
||| > provide correctness-by-construction guarantees that substructural properties of written programs hold.
|||
||| This module defines the Framework's core.
|||
module Resources

%default total
%access public export

namespace TypeLevel

  ||| Language variables are typed.
  |||
  ||| @Ty the type universe in which the variable exists.
  ||| @ty the variable's actual type.
  public export
  data Var : (Ty : Type) -> (ty : Ty) -> Type where
    MkVar : Var type value

  ||| Each variable has an associated abstract state.
  |||
  ||| @type the type universe in which the variable exists.
  ||| @calcStateType a function to compute the state's value from the provided type.
  ||| @value the variables type from which the abstract state's type is calculated.
  public export
  data StateItem : (type : Type)
                -> (calcStateType : type -> Type)
                -> (value : type)
                -> Type
    where
       MkStateItem : (value : type)
                  -> (label : Var type value)
                  -> (state : calcStateType value)
                  -> StateItem type calcStateType value

  ||| The context in which EDSL variables are bound and associated
  ||| with an abstract state.
  |||
  ||| The context is a simple cons-like data structure whose internal
  ||| values are are typed sith a common value.  We cannot utilise
  ||| Idris' `List` type as we need to collect at the type-level the
  ||| variable's type.  In future this should be replaced with `DList`
  ||| from idris-containers but that is future engineering.
  |||
  ||| @type the type universe in which the variable exists.
  ||| @calcStateType a function to compute the state's value from the provided type.
  public export
  data Context : (type : Type)
              -> (calcStateType : type -> Type)
              -> Type
    where
      Nil  : Context type calcStateType
      (::) : (item : StateItem type calcStateType value)
          -> (rest : Context type calcStateType)
          -> Context type calcStateType

  ||| A function to construct the type signature common to all
  ||| /Resource Dependent EDSLs/ as specified in the framework.
  |||
  ||| The signature describes a Hoare like monad that describes the
  ||| shift from an existing `pre` context to a new context based on
  ||| the expression's type.
  |||
  ||| @type the type universe in which the variable exists.
  ||| @calcStateType a function to compute the state's value from the provided type.
  public export
  Lang : (type : Type) -> (calcStateType : (val : type) -> Type) -> Type
  Lang type calcStateType =  (exprTy : Type)
                          -> (pre    : Context type calcStateType)
                          -> (postK  : exprTy -> Context type calcStateType)
                          -> Type

  ||| A simple algebraic data type to represent a /Resource Dependent EDSL/ instance.
  |||
  ||| @type the type universe in which the variable exists.
  ||| @calcStateType a function to compute the state's value from the provided type.
  public export
  data LANG : (type : Type) -> (calcStateType : type -> Type) -> Type where
      MkLang : (type : Type)
            -> (calcStateType : type -> Type)
            -> (lang : Lang type calcStateType)
            -> LANG type calcStateType

  ||| A type-level function calculate the next set of states based on a runtime value.
  |||
  ||| @value The runtime value used for calculation.
  ||| @postK The function that calculates the new state using `value`.
  public export
  updateState : (value : t)
             -> (postK : t -> Context type calcStateType)
             -> Context type calcStateType
  updateState value postK = postK value

  ||| A data structure that captures, generically, the sequencing of resource dependent EDSL language expressions.
  |||
  ||| @m the computation context in which this EDSL is defined.
  ||| @x the type associated with a language expression.
  ||| @calcStateType a function to compute the state's value from the provided type.
  ||| @spec The specific EDSL that is to be embedded.
  ||| @pre  The initial state that the expression knows about.
  ||| @post The final state the expression will be in that is dependent on the runtime value associated with the expression.
  public export
  data LangM : (m : Type -> Type)
            -> (x : Type)
            -> (spec : LANG type calcStateType)
            -> (pre  : Context type calcStateType)
            -> (post : x -> Context type calcStateType)
            -> Type
    where
      ||| Returns a value.
      Value : (value : a)
           -> {postK : a -> Context type calcStateType}
           -> LangM m a spec (postK value) postK

      ||| Provides sequencing of expressions and insertion of computed values into subsequent expressions.
      |||
      ||| @this The computation to run first.
      ||| @beInThis The next computation to run that uses the result of `this`.
      Let : (this : LangM m a spec old oldK)
         -> (beInThis : (val : a) -> LangM m b spec (oldK val) newK)
         -> LangM m b spec old newK

      ||| How to embedd an EDSL language expression into the framework.
      |||
      ||| @t The type of the expression.
      ||| @type the type universe in which the variable exists.
      ||| @calcStateType a function to compute the state's value from the provided type.
      ||| @pre  The initial state that the expression knows about.
      ||| @postK The final state the expression will be in that is dependent on the runtime value associated with the expression.
      ||| @eSig The function to compute the expressions type signature.
      ||| @expr The expression to embedd.
      Expr : {t,type : Type}
          -> {calcStateType : type -> Type}
          -> {pre   : Context type calcStateType}
          -> {postK : t -> Context type calcStateType}
          -> {eSig  : Lang type calcStateType}
          -> (expr : eSig t pre postK)
          -> LangM m t (MkLang type calcStateType eSig) pre (\v => updateState v postK)

  -- A high-level API to make embedding and use of the framework easier for users and designers.
  namespace API
    public export
    pure : (value : a) -> {postK : a -> Context type calcStateType} -> LangM m a spec (postK value) postK
    pure = Value

    public export
    (>>=) : LangM m a spec old oldK
       -> ((val : a) -> LangM m b spec (oldK val) postK)
       -> LangM m b spec old postK
    (>>=) = Let

    public export
    expr : {t,type : Type}
        -> {calcStateType : type -> Type}
        -> {pre   : Context type calcStateType}
        -> {postK : t -> Context type calcStateType}
        -> {eSig  : Lang type calcStateType}
        -> (expr : eSig t pre postK)
        -> LangM m t (MkLang type calcStateType eSig) pre (\v => updateState v postK)
    expr e = Expr e

namespace All
  ||| Ensure that a predicate applies to all items in the given context.
  |||
  ||| @p The predicate to apply to each item.
  ||| @c The context to apply the predicate to.
  public export
  data AllContext : (p : (value : type)
                      -> (item  : StateItem type calcStateType value)
                      -> Type)
                 -> (c : Context type calcStateType)
                 -> Type
    where
      Nil  : AllContext p Nil
      (::) : {p : (v : type) -> StateItem type calcStateType v -> Type}
          -> {value : type}
          -> {item  : StateItem type calcStateType value}
          -> {items : Context type calcStateType}
          -> (prf   : p value item)
          -> (rest  : AllContext p items)
          -> AllContext p (item :: items)

namespace Element

  ||| Existential quantification that there exists an item in the
  ||| context that the given predicate is satisifed.
  |||
  ||| @value The type of variable that the predicate is applied to.
  ||| @p The predicate to apply to each item.
  ||| @c The context to apply the predicate to.
  public export
  data InContext : (value : type)
                -> (p : StateItem type calcStateType value -> Type)
                -> (c : Context type calcStateType)
                -> Type
    where
      H : {p : StateItem type calcStateType value -> Type}
       -> (bewijs : p x)
       -> InContext value p (x :: rest)

      E : {p : StateItem type calcStateType value -> Type}
       -> (komt : InContext value p rest)
       -> InContext value p (not_x :: rest)

  Uninhabited (InContext value p Nil) where
    uninhabited (H _) impossible
    uninhabited (E _) impossible

  ||| Updates the element (at the provided index) in the provided context.
  |||
  ||| @context The context in which to apply the update.
  ||| @index The position of the item to be updated.
  ||| @f The update function to apply. The proof is included incase it contains
  |||    any interesting information for the update.
  public export
  update : (context : Context type calcStateType)
        -> (index   : InContext value predicate context)
        -> (f : (item   : StateItem type calcStateType value)
             -> (bewijs : predicate item)
             -> StateItem type calcStateType value)
        -> Context type calcStateType
  update (item :: rest)     (H bewijs) f = f item bewijs :: rest
  update (not_item :: rest) (E komt)   f = not_item :: update rest komt f

  ||| Remove the element (at the provided index) in the provided context.
  |||
  ||| @predicate The predicate that was applied to the item to be dropped.
  ||| @context The context in which the item is to be dropped.
  ||| @index The position of the item to drop.
  public export
  drop : {predicate : StateItem type calcStateType value -> Type}
      -> (context   : Context type calcStateType)
      -> (index     : InContext value predicate context)
      -> Context type calcStateType
  drop (item :: rest) (H bewijs) = rest
  drop (not_x :: rest) (E komt) = not_x :: drop rest komt

  ||| A helper function to replace the state in an element in our context.
  public export
  doUpdate : {predicate : StateItem type calcStateType value -> Type}
          -> (item     : StateItem type calcStateType value)
          -> (newState : calcStateType value)
          -> (bewijs   : predicate item)
          -> StateItem type calcStateType value
  doUpdate (MkStateItem value label state) newState bewijs = (MkStateItem value label newState)

  ||| An application of `update` to set the state for an item at the specified index.
  |||
  ||| @predicate The predicate that was applied to the item.
  ||| @context The context in which to apply the update.
  ||| @index The position of the item to be updated.
  ||| @item' The new state.
  public export
  setState : {predicate : StateItem type calcStateType value -> Type}
          -> (context   : Context type calcStateType)
          -> (index     : InContext value predicate context)
          -> (item'     : calcStateType value)
          -> Context type calcStateType
  setState context index item' = update context index (\i => doUpdate i item')

  ||| Proof that the given variable has the associated state.
  data HasVar : (lbl  : Var type value)
             -> (item : StateItem type calcStateType value)
             -> Type where
    DoesHaveVar : HasVar lbl (MkStateItem value lbl state)

-- -------------------------------------------------------------- [ Evaluation ]

namespace Env

  ||| Captures how EDSL variables are translated into concrete types.
  public export
  interface RealVar (type : Type) where
    CalcRealType : type -> Type

  ||| A container for holding concrete variable representations.
  public export
  data Tag : (type : Type) -> (value : type)
          -> Type
    where
      MkTag : RealVar type
           => (real : CalcRealType value)
           -> Tag type value

  ||| An evaluation environments to keep track of our language's real
  ||| variables and their abstract state, respective to the language's
  ||| typing context.
  public export
  data Env : (m : Type -> Type)
          -> (type    : Type)
          -> {calcStateType : type -> Type}
          -> (ctxt    : Context type calcStateType)
          -> Type
    where
      Nil  : Env m type Nil
      (::) : RealVar type
          => {item : StateItem type calcStateTy value}
          -> (tag  : Tag type value)
          -> (rest : Env m type items)
          -> Env m type (item::items)

  ||| Extracting a variable from the environment at the given index.
  |||
  ||| @env The environment being accessed.
  ||| @idx The position of the variable value to be extracted.
  public export
  lookup : RealVar type
        => (env : Env m type context)
        -> {p   : (item : StateItem type calcSTy value) -> Type}
        -> (idx : InContext value p context)
        -> Tag type value
  lookup (tag :: rest) (H bewijs) = tag
  lookup (tag :: rest) (E komt) = lookup rest komt

  ||| Update the type level state in the environment.
  |||
  ||| @env The environment being accessed.
  ||| @idx The position of the variable value to be updated
  ||| @up  The update function to be applied.
  public export
  update : RealVar type
        => (env  : Env m type ctxt)
        -> (idx  : InContext value p ctxt)
        -> (up   : (i : StateItem type calcSTy value)
                -> p i
                -> StateItem type calcSTy value)
        -> Env m type (update ctxt idx up)
  update (tag :: rest) (H bewijs) up = tag :: rest
  update (tag :: rest) (E komt) up = tag :: update rest komt up

  ||| Drop the state at the given index.
  |||
  ||| @env The environment being accessed.
  ||| @idx The position of the variable value to be dropped.
  public export
  drop : (env : Env m type ctxt)
      -> (idx : InContext value p ctxt)
      -> Env m type (drop ctxt idx)
  drop (tag :: rest) (H bewijs) = rest
  drop (tag :: rest) (E komt) = tag :: drop rest komt

namespace Building

  ||| The handler interface that describes how to evaluate languages in a presented computation context.
  |||
  ||| @type the type universe in which the variable exists.
  ||| @calcStateType a function to compute the state's value from the provided type.
  ||| @eSig The function to compute the expressions type signature.
  ||| @tyBRes The type associated with the accumulator if we are to build something.
  ||| @m The computation context in which to evaluate the language.
  public export
  interface RealVar type
         => Handler (type : Type)
                    (calcStateType : type -> Type)
                    (eSig   : Lang type calcStateType)
                    (tyBRes : Type)
                    (m      : Type -> Type) | type
    where
      ||| How to handle language expressions in the given computation context `m`.
      |||
      ||| @env the current evaluation context.
      ||| @expr the expression to be handled.
      ||| @acc an accumulator to be used if building something.
      ||| @cont a continuation to make things total.
      handle : (env  : Env m type pre)
            -> (expr : eSig tyExpr pre postK)
            -> (acc  : tyBRes)
            -> (cont : (value : tyExpr)
                    -> (env'  : Env m type (postK value))
                    -> (acc'  : tyBRes)
                    -> m tyRes)
            -> m tyRes

  ||| An internal helper function to evaluate an embedded language expression.
  |||
  ||| You will not need to call this at all.
  public export
  execBuild : RealVar type
           => Handler type calcStateType eSig tyBRes m
           => (env  : Env m type pre)
           -> (expr : eSig tyExpr pre postK)
           -> (acc  : tyBRes)
           -> (cont : (value : tyExpr)
                   -> (env'  : Env m type (updateState value postK))
                   -> (acc'  : tyBRes)
                   -> m tyRes)
           -> m tyRes
  execBuild env expr acc cont = handle env expr acc cont

  ||| An internal helper function to evaluate a `LangM` instance, and sequencing of expressions.
  |||
  ||| You will not need to call this at all.
  public export
  rawEval : RealVar type
         => Handler type calcStateType eSig tyBRes m
         => (env  : Env m type pre)
         -> (expr : LangM m tyExpr (MkLang type calcStateType eSig) pre postK)
         -> (acc  : tyBRes)
         -> (cont : (value : tyExpr)
                 -> (env'  : Env m type (postK value))
                 -> (acc'  : tyBRes)
                 -> m tyRes)
         -> m tyRes
  rawEval env (Value value {postK}) acc cont = cont value env acc
  rawEval env (Let this beIn) acc cont =
    rawEval env this acc (\this', env', acc' =>
                             rawEval env' (beIn this') acc' cont)
  rawEval env (Expr expr) acc cont = execBuild env expr acc cont

  ||| A helper data structure to container evaluation results.
  public export
  record Result (tyExpr : Type) (tyBRes : Type) where
    constructor MkResult
    computed : tyExpr
    built    : tyBRes

  ||| Evaluate the given closed program in the given applicative computation context `m`.
  |||
  ||| @init the initial value for our accumulator.
  ||| @expr the program to evaluate.
  public export
  run : Applicative m
     => RealVar type
     => Handler type calcStateType eSig tyBRes m
     => (init : tyBRes)
     -> (expr : LangM m tyExpr (MkLang type calcStateType eSig) Nil (const Nil))
     -> m (Result tyExpr tyBRes)
  run init expr = rawEval Nil expr init (\val, env', acc' => pure $ MkResult val acc')

  ||| Evaluate the given closed program in the identity computation context.
  |||
  ||| @init the initial value for our accumulator.
  ||| @expr the program to evaluate.
  public export
  runPure : RealVar type
        => Handler type calcStateType eSig tyBRes Basics.id
        => (init : tyBRes)
        -> (expr : LangM Basics.id tyExpr (MkLang type calcStateType eSig) Nil (const Nil))
        -> (Result tyExpr tyBRes)
  runPure init expr = rawEval Nil expr init (\val, env', acc' => MkResult val acc')

-- --------------------------------------------------------------------- [ EOF ]


module ATM

import Data.Vect

PIN: Type
PIN = Vect 4 Integer

data ATMState = Ready | CardInserted | Session

data HasCard : ATMState -> Type where
  HasCI : HasCard CardInserted
  HasSession : HasCard Session

data PINCheck = CorrectPIN | IncorrectPIN

data ATMCmd : (ty: Type) -> ATMState -> (ty -> ATMState) -> Type where
  InsertCard: ATMCmd () Ready (const CardInserted)
  EjectCard: {auto prf: HasCard state} -> ATMCmd () state (const Ready)
  GetPIN: ATMCmd PIN CardInserted (const CardInserted)
  CheckPIN: PIN -> ATMCmd PINCheck CardInserted (\check => case check of
                                                              CorrectPIN => Session
                                                              IncorrectPIN => CardInserted)

  GetAmount: ATMCmd Nat state (const state)
  Dispense: (amount: Nat) -> ATMCmd () Session (const Session)

  Message: String -> ATMCmd () state (const state)

  Pure: (res: ty) -> ATMCmd ty (state_fn res) state_fn

  (>>=): ATMCmd a state1 state2_fn -> ((res: a) -> ATMCmd b (state2_fn res) state3_fn) -> ATMCmd b state1 state3_fn

insertEject : ATMCmd () Ready (const Ready)
insertEject = do InsertCard
                 EjectCard

-- atm: ATMCmd () Ready (const Ready)
-- atm = do InsertCard
--          pin <- GetPIN
--          pinOK <- CheckPIN pin
--          case pinOK of
--               CorrectPIN => do cash <- GetAmount
--                                Dispense cash
--                                EjectCard
--               IncorrectPIN => EjectCard

-- runATM: ATMCmd res inState outState_fn -> IO res
-- runATM InsertCard = do putStrLn "Please insert your card"
--                        x <- getLine
--                        pure()
-- runATM EjectCard = putStrLn "Thank you"
-- runATM GetPIN = do putStr "Enter Pin"
--                    pure (the (Vect _ Integer) [1, 2, 3, 4])
-- runATM (CheckPIN pin) = if pin == [1, 2, 3, 4]
--                           then pure CorrectPIN
--                           else pure IncorrectPIN
-- runATM GetAmount = do putStr "How much would you like"
--                       x <- getLine
--                       pure (cast x)
-- runATM (Dispense amount) = putStrLn "Here is the amount"
-- runATM (Message msg) = putStrLn msg
-- runATM (Pure res) = pure res
-- runATM (x >>= f) = do res <- runATM x
--                       runATM (f res)


-- This regtest tests for incorrect erasure from underapplied data constructors.

data T : Type where
    C : Nat -> Nat -> Nat -> T

data Q : Type where
    D : (Nat -> T) -> Q

-- The bug was here: (C 3 4) would compile to plain data,
-- not to a function that would expect another (albeit erased) argument.
f : Q
f = D (C 3 4)

proj : Nat -> Q -> Nat -> T
proj    Z  (D f) = f
proj (S n)    q  = proj n q

g : T -> Nat
g (C x y z) = x + y

-- Here, we'd apply the not-a-function to the erased argument 4,
-- which would make the program go wrong.
main : IO ()
main = printLn $ g (proj 3 f 4)


import Data.Vect

fourInts: Vect 4 Int
fourInts = [1, 2, 2, 4]

sixInts: Vect 6 Int
sixInts = [1, 2, 3, 4, 5, 7]

tenInts: Vect 10 Int
tenInts = fourInts ++ sixInts


module Data.Cotree

import Data.Colist
import Data.Maybe
import Data.Tree

%default total

--------------------------------------------------------------------------------
--          Cotrees: Potentially infinte trees
--------------------------------------------------------------------------------

mutual
  ||| A potentially infinite rose tree
  public export
  record Cotree (a : Type) where
    constructor MkCotree
    value  : a
    forest : Inf (Coforest a)

  ||| A potentially finit stream of trees
  public export
  Coforest : Type -> Type
  Coforest = Colist . Cotree

public export
singleton : a -> Cotree a
singleton a = MkCotree a Nil

public export
unfold : (f : b -> (a,Colist b)) -> b -> Cotree a
unfold f vb = let (va,bs) = f vb
                   in MkCotree va $ unfoldF bs

  where unfoldF : Colist b -> Coforest a
        unfoldF []       = []
        unfoldF (h :: t) = unfold f h :: unfoldF t

public export
iterate : (f : a -> Colist a) -> a -> Cotree a
iterate f a = unfold (\v => (v, f v)) a

public export
expand : (a -> Colist a) -> Cotree a -> Cotree a
expand f (MkCotree v vs) = let MkCotree v2 vs2 = iterate f v
                            in MkCotree v2 (run vs vs2)
  where run : Coforest a -> Coforest a -> Coforest a
        run []        ys = ys
        run (x :: xs) ys = expand f x :: run xs ys

--------------------------------------------------------------------------------
--          To and from Tree
--------------------------------------------------------------------------------

public export
fromTree : Tree a -> Cotree a
fromTree (MkTree v fo) = MkCotree v (fromForest fo)
  where fromForest : Forest a -> Coforest a
        fromForest []        = []
        fromForest (x :: xs) = fromTree x :: fromForest xs

||| Converts a Cotree to a tree of the given maximum depth and width.
||| The maximum numbers of elements in the tree will be
||| maxWidth ^ maxDepth.
public export
toTree : (maxDepth : Nat) -> (maxWidth : Nat) -> Cotree a -> Tree a
toTree 0     _  (MkCotree v fo) = MkTree v []
toTree (S k) mw (MkCotree v fo) = MkTree v (toForest mw fo)
  where toForest : Nat -> Coforest a -> Forest a
        toForest 0     _         = []
        toForest (S n) []        = []
        toForest (S n) (t :: ts) = toTree k mw t :: toForest n ts

--------------------------------------------------------------------------------
--          Functor and Applicative
--------------------------------------------------------------------------------

public export
mapCotree : (a -> b) -> Cotree a -> Cotree b
mapCotree f (MkCotree v vs) = MkCotree (f v) (mapForest vs)
  where mapForest : Coforest a -> Coforest b
        mapForest []       = []
        mapForest (h :: t) = mapCotree f h :: mapForest t

public export
interleave : Cotree (a -> b) -> Cotree a -> Cotree b
interleave tf@(MkCotree vf fs) ta@(MkCotree va as) =
  MkCotree (vf va) (interleaveFs fs)

  where interleaveAs : Coforest a -> Coforest b
        interleaveAs []       = []
        interleaveAs (h :: t) = interleave tf h :: interleaveAs t

        interleaveFs : Coforest (a -> b) -> Coforest b
        interleaveFs []       = interleaveAs as
        interleaveFs (h :: t) = interleave h ta :: interleaveFs t

public export
bind : Cotree a -> (a -> Cotree b) -> Cotree b
bind (MkCotree v vs) f = let MkCotree w ws = f v
                       in MkCotree w (run vs ws)
  where run : Coforest a -> Coforest b -> Coforest b
        run []        ys = ys
        run (x :: xs) ys = bind x f :: run xs ys

public export
bindMaybe : Cotree (Maybe a) -> (a -> Cotree (Maybe b)) -> Cotree (Maybe b)
bindMaybe (MkCotree mv tas) f =
  case map f mv of
       Nothing => MkCotree Nothing (run tas Nil)
       Just (MkCotree mb tbs) => MkCotree mb (run tas tbs)

  where run : Coforest (Maybe a) -> Coforest (Maybe b) -> Coforest (Maybe b)
        run [] ys        = ys
        run (x :: xs) ys = bindMaybe x f :: run xs ys

--------------------------------------------------------------------------------
--          Shrinking
--------------------------------------------------------------------------------

public export
shrink : (maxSteps : Nat) -> Cotree (Maybe a) -> List a
shrink maxSteps x = run maxSteps [x]
  where run : Nat -> Coforest (Maybe a) -> List a
        run _ Nil          = Nil
        run 0 _            = Nil
        run (S k) (h :: t) = case h.value of
                               Just a  => a :: run k h.forest
                               Nothing => run k t

public export
mapShrink : (maxSteps : Nat) -> (a -> Maybe b) -> Cotree a -> List b
mapShrink ms f = shrink ms . mapCotree f

public export
shrinkIf : (maxSteps : Nat) -> (a -> Bool) -> Cotree a -> List a
shrinkIf ms p = mapShrink ms (\a => if p a then Just a else Nothing)

||| Prunes a cotree up to the given depth and width.
public export
pruneTo : (width : Nat) -> (depth : Nat) -> Cotree a -> Cotree a
pruneTo _ 0     (MkCotree v _ ) = MkCotree v Nil
pruneTo w (S d) (MkCotree v vs) = MkCotree v $ (map (pruneTo w d) $ keep w vs)
  where keep : Nat -> Colist t -> Colist t
        keep _ [] = []
        keep 0 _  = []
        keep (S k) (x :: xs) = x :: keep k xs

||| Removes all children from a cotree
public export
prune : Cotree a -> Cotree a
prune = pruneTo 0 0

--------------------------------------------------------------------------------
--          Interfaces
--------------------------------------------------------------------------------

public export
Functor Cotree where
  map = mapCotree

public export
Applicative Cotree where
  pure  = singleton
  (<*>) = interleave

--------------------------------------------------------------------------------
--          Filtering
--------------------------------------------------------------------------------

public export
takeUntil : (a -> Bool) -> Cotree a -> Cotree a
takeUntil f (MkCotree v vs) = if f v then MkCotree v []
                                     else MkCotree v (takeUntilF vs)
  where takeUntilF : Coforest a -> Coforest a
        takeUntilF []        = vs
        takeUntilF (MkCotree x xs :: ts) =
          if f x
             then [MkCotree x []]
             else MkCotree x (takeUntilF xs) :: takeUntilF ts

public export
takeBeforeNothing : Cotree (Maybe a) -> Maybe (Cotree a)
takeBeforeNothing (MkCotree Nothing _) = Nothing
takeBeforeNothing (MkCotree (Just v) vs) = Just (MkCotree v (run vs))
  where run : Coforest (Maybe a) -> Coforest a
        run []                             = []
        run ((MkCotree Nothing _) :: _)    = []
        run ((MkCotree (Just v) vs) :: ts) = MkCotree v (run vs) :: run ts

public export
takeBefore : (a -> Bool) -> Cotree a -> Maybe (Cotree a)
takeBefore f = takeBeforeNothing . map (\a => toMaybe (f a) a)

public export %inline
takeWhile : (a -> Bool) -> Cotree a -> Maybe (Cotree a)
takeWhile f = takeBefore (not . f)

public export %inline
mapMaybe : (a -> Maybe b) -> Cotree a -> Maybe (Cotree b)
mapMaybe f = takeBeforeNothing . map f


module CH06.DataStore

import Data.Vect

infixr 5 .+.

data Schema = SString
            | SInt
            | SChar
            | (.+.) Schema Schema

SchemaType : Schema -> Type
SchemaType SString   = String
SchemaType SInt      = Int
SchemaType SChar     = Char
SchemaType (x .+. y) = (SchemaType x, SchemaType y)

data DataStore : Type where
  MkData : (schema : Schema)
         -> (size : Nat)
         -> (items : Vect size (SchemaType schema))
         -> DataStore

data Command : Schema -> Type where
  SetSchema : (newschema : Schema)  -> Command schema
  Add  : SchemaType schema -> Command schema
  Get  : Integer -> Command schema
  Size : Command schema
  Quit : Command schema

display : SchemaType schema -> String
display {schema = SString} item             = item
display {schema = SInt} item                = show item
display {schema = SChar} item               = show item
display {schema = (x .+. y)} (iteml, itemr) = display iteml ++ ", " ++ display itemr

parsePrefix : (schema : Schema) -> String -> Maybe (SchemaType schema, String)
parsePrefix SChar input                 = case unpack input of
                                            [x]              => Just (x, "")
                                            (x :: ' ' :: xs) => Just (x, ltrim $ pack xs)
                                            _                => Nothing
parsePrefix SString input               = getQuoted (unpack input)
  where
    getQuoted : List Char -> Maybe (String, String)
    getQuoted ('"' :: xs) =
      case span (/= '"') xs of
        (quoted, '"' :: rest) => Just (pack quoted, ltrim (pack rest))
        _                     => Nothing
    getQuoted _           = Nothing
parsePrefix SInt input                  = case span isDigit input of
                                          ("", _)     => Nothing
                                          (num, rest) => Just (cast num, ltrim rest)
parsePrefix (schemal .+. schemar) input = do
  (lval, inputl) <- parsePrefix schemal input
  (rval, inputr) <- parsePrefix schemar inputl
  pure ((lval, rval), inputr)

parseBySchema : (schema : Schema) -> String -> Maybe (SchemaType schema)
parseBySchema schema input =
  case parsePrefix schema input of
    Just (res, "") => Just res
    _              => Nothing

size : DataStore -> Nat
size (MkData schema' size' items') = size'

schema : DataStore -> Schema
schema (MkData schema' size' items') = schema'

items : (store : DataStore) -> Vect (size store) (SchemaType (schema store))
items (MkData schema' size' items') = items'

addToStore : (store : DataStore) -> SchemaType (schema store) -> DataStore
addToStore (MkData schema' size' store') newitem = MkData schema' _ (addToData store')
  where
    addToData : Vect oldSize (SchemaType schema') -> Vect (S oldSize) (SchemaType schema')
    addToData []        = [newitem]
    addToData (x :: xs) = x :: addToData xs

setSchema : (store : DataStore) -> Schema -> Maybe DataStore
setSchema store schema = case size store of
                           Z => Just (MkData schema _ [])
                           _ => Nothing

parseSchema : List String -> Maybe Schema
parseSchema ("String" :: xs) =
  case xs of
    [] => Just SString
    _  => do
      schema1 <- parseSchema xs
      Just (SString .+. schema1)
parseSchema ("Char" :: xs)   =
  case xs of
    [] => Just SChar
    _  => do
      schema1 <- parseSchema xs
      Just (SChar .+. schema1)
parseSchema ("Int" :: xs)    =
  case xs of
    [] => Just SInt
    _  => do
      schema1 <- parseSchema xs
      Just (SInt .+. schema1)
parseSchema _                = Nothing

parseCommand : (schema : Schema) -> (cmd : String) -> (args : String) -> Maybe (Command schema)
parseCommand schema "schema" rest = SetSchema <$> (parseSchema (words rest))
parseCommand schema "add" val     = Add <$> (parseBySchema schema val)
parseCommand schema "get" val     = if all isDigit (unpack val)
                                    then Nothing
                                    else Just (Get (cast val))
parseCommand schema "quit" ""     = Just Quit
parseCommand schema "size" ""     = Just Size
parseCommand _ _ _                = Nothing

parse : (schema : Schema) -> (input : String) -> Maybe (Command schema)
parse schema input = case span (/= ' ') input of
                       (cmd, args) => parseCommand schema cmd (ltrim args)

getEntry : (pos : Integer) -> (store : DataStore) -> Maybe (String, DataStore)
getEntry pos store = let store_items = items store
                     in case integerToFin pos (size store) of
                          Nothing => Just ("Out of range\n", store)
                          Just id => Just ((display $ index id store_items) ++ "\n", store)

processInput : DataStore -> String -> Maybe (String, DataStore)
processInput store input =
  case parse (schema store) input of
    Nothing                  => Just ("Invalid command\n", store)
    Just (SetSchema schema') => case setSchema store schema' of
                                  Nothing     => Just ("Can't update scema\n", store)
                                  Just store' => Just ("OK\n", store')
    Just (Add item)          => Just ("ID " ++ show (size store) ++ "\n", addToStore store item)
    Just (Get pos)           => getEntry pos store
    Just Size                => Just("Size " ++ show (size store) ++ "\n", store)
    Just Quit                => Nothing

main : IO ()
main = replWith (MkData (SString .+. SString .+. SInt) _ []) "Command: " processInput


||| Data type `Validated` was introduced in chapter
||| *Functor and Friends*. It provides an error type
||| with the ability of error accumulation.
module Data.Validated

%default total

public export
data Validated : (e,a : Type) -> Type where
  Invalid : (err : e) -> Validated e a
  Valid   : (val : a) -> Validated e a

--------------------------------------------------------------------------------
--          Utilities
--------------------------------------------------------------------------------

public export
validated : (e -> b) -> (a -> b) -> Validated e a -> b
validated f _ (Invalid ve) = f ve
validated _ g (Valid va)   = g va

public export
fromEither : Either e a -> Validated e a
fromEither = either Invalid Valid

public export
toEither : Validated e a -> Either e a
toEither = validated Left Right

--------------------------------------------------------------------------------
--          Interfaces
--------------------------------------------------------------------------------

public export
Eq e => Eq a => Eq (Validated e a) where
  Invalid x == Invalid y = x == y
  Valid x   == Valid y   = x == y
  _         == _         = False

public export
Ord e => Ord a => Ord (Validated e a) where
  compare (Invalid x) (Invalid y) = compare x y
  compare (Invalid _) (Valid _)   = LT
  compare (Valid _)   (Invalid _) = GT
  compare (Valid x)   (Valid y)   = compare x y

public export
Show e => Show a => Show (Validated e a) where
  showPrec p (Invalid e) = showCon p "Invalid" (showArg e)
  showPrec p (Valid a)   = showCon p "Valid" (showArg a)

public export
Functor (Validated e) where
  map _ (Invalid err) = Invalid err
  map f (Valid val)   = Valid $ f val

public export
Semigroup e => Applicative (Validated e) where
  pure = Valid
  Valid   f  <*> Valid v    = Valid $ f v
  Valid   _  <*> Invalid ve = Invalid ve
  Invalid e1 <*> Invalid e2 = Invalid $ e1 <+> e2
  Invalid ve <*> Valid _    = Invalid ve

public export
Foldable (Validated e) where
  foldr acc st (Invalid _) = st
  foldr acc st (Valid v)   = acc v st

  foldl acc st (Invalid _) = st
  foldl acc st (Valid v)   = acc st v

  foldlM acc st (Invalid _) = pure st
  foldlM acc st (Valid v)   = acc st v

  foldMap f (Invalid _) = neutral
  foldMap f (Valid v)   = f v

  toList (Invalid _) = []
  toList (Valid v)   = [v]

  null (Invalid _) = True
  null (Valid v)   = False

public export
Traversable (Validated e) where
  traverse f (Invalid ve) = pure $ Invalid ve
  traverse f (Valid va)   = Valid <$> f va

public export
Bifunctor Validated where
  bimap f g (Invalid ve) = Invalid $ f ve
  bimap f g (Valid va)   = Valid $ g va

  mapFst f (Invalid ve) = Invalid $ f ve
  mapFst f (Valid va)   = Valid va

  mapSnd f (Invalid ve) = Invalid ve
  mapSnd f (Valid va)   = Valid $ f va

public export
Bifoldable Validated where
  bifoldr f g st (Invalid ve) = f ve st
  bifoldr f g st (Valid va)   = g va st

  bifoldl f g st (Invalid ve) = f st ve
  bifoldl f g st (Valid va)   = g st va

  binull _ = False

public export
Bitraversable Validated where
  bitraverse f g (Invalid ve) = Invalid <$> f ve
  bitraverse f g (Valid va)   = Valid <$> g va

public export
Monoid e => Alternative (Validated e) where
  empty = Invalid neutral

  Valid va  <|> _         = Valid va
  Invalid x <|> Invalid y = Invalid $ x <+> y
  Invalid _ <|> Valid va  = Valid va