Polynomial lens in Idris2

PolyLens.idr

module Main

data Vect : Type -> Nat -> Type where
  Nil : Vect a Z
  (::) : (x: a) -> (xs : Vect a n) -> Vect a (S n)

data Tree : Type -> Nat -> Type where
  Empty : Tree a Z
  Leaf  : a -> Tree a (S Z)
  Node  : Tree a k1 -> Tree a k2 -> Tree a (S (plus k1  k2))

-- Synonym for Nat-dependent types
Nt : Type
Nt = Nat -> Type

-- Existential types hiding dependence on Nat
data Some : (nt : Nat -> Type) -> Type where
  Hide : {0 n : Nat} -> nt n -> Some nt

SomeVect : Type -> Type
SomeVect a = Some (Vect a)

SomeTree : Type -> Type
SomeTree a = Some (Tree a)

-- Vector utility functions

mapV : (a -> b) -> Vect a n -> Vect b n
mapV f Nil = Nil
mapV f (a :: v) = (f a) :: (mapV f v)

concatV : Vect a m -> Vect a n -> Vect a (plus m n)
concatV Nil v = v
concatV (a :: w) v = a :: (concatV w v)

splitV : (n : Nat) -> Vect a (plus n m) -> (Vect a n, Vect a m)
splitV Z v = (Nil, v)
splitV (S k) (a :: v') = let (v1, v2) = splitV k v'
                            in (a :: v1, v2)

sortV : Ord a => Vect a n -> Vect a n
sortV Nil = Nil
sortV (x :: xs) = let xsrt = sortV xs 
                  in (ins x xsrt)
  where -- a is the same type as in outer scope, but we use a new n'
    ins : a -> Vect a n' -> Vect a (S n')
    ins x Nil = x :: Nil
    ins x (y :: xs) = if x < y 
                      then x :: (y :: xs)
                      else y :: ins x xs

-- The lens returns this existential type
-- For instance:
-- exists n. (k, Vect a n, Vect b n -> Tree k)
-- is a pair: a vector of leaves and a leaf changer
data SomePair : Nt -> Nt -> Nt -> Type where
 HidePair : {0 n : Nat} -> (k : Nat) -> a n -> (b n -> t k) -> SomePair a b t

------------------
-- Polynomial Lens
------------------

PolyLens : Nt -> Nt -> Nt -> Nt -> Type
PolyLens s t a b = {k : Nat} -> s k -> SomePair a b t

getLens :  PolyLens s t a b -> {k : Nat} -> s k -> Some a
getLens lens sk = 
  let (HidePair k' an f) = lens sk
  in Hide an

transLens : PolyLens s t a b -> 
   (forall m. a m -> b m) -> {k : Nat} -> s k -> Some t
transLens lens f sk =
  let (HidePair k an bn_tk) : SomePair a b t = lens sk
  in  Hide (bn_tk (f an))

--------------------
-- Tree Lens
-- focuses on leaves
--------------------
-- When traversing the tree, each branch produces
-- a leaf changer. We have to compose them

compose : {n : Nat} -> --{m : Nat} -> {k1 : Nat} -> {k2 : Nat} ->
       (Vect b n -> Tree b k1) -> (Vect b m -> Tree b k2) ->
       Vect b (n + m) -> Tree b (S(k1 + k2))
compose {n} {m} f1 f2 v =
  let (v1, v2) = splitV n v
  in Node (f1 v1) (f2 v2)



-- Given source Tree a, return a pair
--  (Vect a of leaves, function from Vect b of leaves to Tree b)
replace : (b : Type) -> Tree a k -> (n : Nat ** (Vect a n, Vect b n -> Tree b k))
replace b Empty = (Z ** (Nil, \v => Empty))
replace b (Leaf x) = (1 ** (x :: Nil, \(y :: Nil) => Leaf y))
replace b (Node t1 t2) with (replace b t1)
  _ | (n1 ** (v1, f1)) with (replace b t2)
    _ | (n2 ** (v2, f2)) =
         let f3 = compose f1 f2
         in (n1 + n2 ** (concatV v1 v2, f3))

treeLens : {a : Type} -> {b : Type} -> PolyLens (Tree a) (Tree b) (Vect a) (Vect b)
treeLens t with (replace b t)
  _ | (n ** (v, f)) = HidePair k v f

treeSimpleLens : {a : Type} -> PolyLens (Tree a) (Tree a) (Vect a) (Vect a)
treeSimpleLens t  with (replace a t)
  _ | (n ** (v, f)) = HidePair k v f

-- Use tree lens to extract a vector of leaves

-- getLens :  PolyLens s t a b -> {k : Nat} -> s k -> Some a
-- treeSimpleLens : {a : Type} -> PolyLens (Tree a) (Tree a) (Vect a) (Vect a)
getLeaves : {a : Type} -> {k : Nat} -> Tree a k -> Some (Vect a)
getLeaves {k} tr = 
  -- this doesn't work 
  -- getLens treeSimpleLens tr
  let (HidePair k' an f) = treeSimpleLens tr
  in Hide an

-- Use tree lens to modify leaves
-- 1. extract the leaves
-- 2. map function over vector of leaves
-- 3. replace leaves with the new vector
mapLeaves : {a : Type} -> {b : Type} -> {k : Nat} -> 
    (a -> b) -> Tree a k -> SomeTree b
mapLeaves {a} {b} f t = 
  let  (HidePair k' v vt) = treeLens t
  in  Hide (vt (mapV f v))

trLeaves : {a : Type} -> {b : Type} -> {n : Nat} -> 
    (forall m. Vect a m -> Vect b m) -> Tree a n -> SomeTree b
trLeaves {a} {b} f tr = 
  -- this doesn't work
  -- transLens treeLens f tr
  let (HidePair k an bn_tk) = treeLens tr
  in  Hide (bn_tk (f an))


-- Utility functions for testing

Show a => Show (Vect a n) where
  show Nil = ""
  show (a :: v) = show a ++ show v

Show a => Show (Tree a n) where
  show Empty = "()"
  show (Leaf a) = show a
  show (Node t1 t2) = "(" ++ show t1 ++ "," ++ show t2 ++ ")"

v0 : Vect Char 0
v0 = Nil
v1 : Vect Char 1
v1 = 'a' :: Nil
v2 : Vect Char 2
v2 = 'b' :: ('c' :: Nil)
v3 : Vect Char 3
v3 = concatV v1 v2

someV : SomeVect Char
someV = Hide v2

Show a => Show (SomeVect a) where
  show (Hide v) = show v

Show a => Show (SomeTree a) where
  show (Hide v) = show v

t3 : Tree Char 5
t3 = (Node (Leaf 'z') (Node (Leaf 'a') (Leaf 'b')))

main : IO ()
main = do
  putStrLn "getLeaves"
  print (getLeaves t3)
  putStrLn "\nmapLeaves"
  print (mapLeaves ord t3)
  putStrLn "\nsortLeaves"
  print (trLeaves sortV t3)
  putStrLn "\nmapLeaves of an empty tree"
  print (mapLeaves ord Empty)
  putStrLn ""