module Soat.Data.Sublist

import public Control.Relation
import Data.List.Elem
import Soat.Data.Product

%default total

infix 5 ++

public export
data Sublist : Rel (List a) where
  Nil  : Sublist [] ys
  (::) : (mem : Elem x ys) -> (mems : Sublist xs ys) -> Sublist (x :: xs) ys

public export
curry : xs `Sublist` ys -> forall x . Elem x xs -> Elem x ys
curry (mem :: mems) Here      = mem
curry (mem :: mems) (There y) = curry mems y

public export
uncurry : {xs : List a} -> (forall x . Elem x xs -> Elem x ys) -> xs `Sublist` ys
uncurry {xs = []}        f = []
uncurry {xs = (x :: xs)} f = f Here :: uncurry (f . There)

public export
uncurryCurry : (f : xs `Sublist` ys) -> uncurry (curry f) = f
uncurryCurry []            = Refl
uncurryCurry (mem :: mems) = cong ((::) mem) (uncurryCurry mems)

public export
curryUncurry : (0 f : forall x . Elem x xs -> Elem x ys) -> (i : Elem x xs)
  -> curry (uncurry f) i = f i
curryUncurry f Here         = Refl
curryUncurry f (There elem) = curryUncurry (f . There) elem

public export
Reflexive (List a) Sublist where
  reflexive = uncurry id

public export
Transitive (List a) Sublist where
  transitive f g = uncurry (curry g . curry f)

public export
elemJoinL : Elem x xs -> Elem x (xs ++ ys)
elemJoinL Here         = Here
elemJoinL (There elem) = There (elemJoinL elem)

public export
elemJoinR : (xs : List a) -> Elem x ys -> Elem x (xs ++ ys)
elemJoinR []        elem = elem
elemJoinR (x :: xs) elem = There (elemJoinR xs elem)

public export
elemSplit : (xs : List a) -> Elem x (xs ++ ys) -> Either (Elem x xs) (Elem x ys)
elemSplit []        elem         = Right elem
elemSplit (x :: xs) Here         = Left Here
elemSplit (x :: xs) (There elem) = bimap There id (elemSplit xs elem)

public export
(++) : {xs, ys, us : List a}
  -> xs `Sublist` ys -> us `Sublist` vs -> (xs ++ us) `Sublist` (ys ++ vs)
(++) mems mems' = uncurry
  (either (elemJoinL . curry mems) (elemJoinR ys . curry mems') . elemSplit xs)

public export
shuffle : is `Sublist` js -> x ^ js -> x ^ is
shuffle []            xs = []
shuffle (mem :: mems) xs = index xs mem :: shuffle mems xs