An implementation and extension of ideas from "Weaving a Web" by Hinze and Jeuring.

Weave.hs

{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}

module Weave
  ( Weaver
  , cons
  , Loc(fdown, fleft, fright, fup, it)
  , down
  , up
  , left
  , right
  , explore
  ) where

import           Control.Monad.ST               ( ST )
import           Data.Function                  ( (&) )
import qualified Data.List.NonEmpty            as NE
import qualified Data.Vector                   as V
import           Data.Vector                    ( (!) )
import qualified Data.Vector.Mutable           as M

newtype Weaver a = W { unW :: (a -> Loc a) -> Loc a }

newtype STWeaver s a = WS { unWS :: (a -> STLoc s a) -> STLoc s a }

callM :: (a -> STWeaver s a) -> (a -> STLoc s a) -> a -> STLoc s a
callM wv fl0 t = unWS (wv t) fl0

call :: (a -> Weaver a) -> (a -> Loc a) -> a -> Loc a
call wv fl0 t = unW (wv t) fl0

-- consM :: (a -> STWeaver s a) -> (M.STVector s a -> a) -> (M.STVector s a -> STWeaver s a)
-- consM wv k ts = WS (\fl0 -> locsM (callM wv) (\ts -> fl0 (k ts)) ts)

consV :: (a -> Weaver a) -> (V.Vector a -> a) -> (V.Vector a -> Weaver a)
consV wv k ts = W (\fl0 -> locsV (call wv) (fl0 . k) ts)

cons :: (a -> Weaver a) -> ([a] -> a) -> ([a] -> Weaver a)
cons wv k ts = W (\fl0 -> locs (call wv) (fl0 . k) ts)

locsM
  :: forall a s
   . (forall s . (a -> ST s (STLoc s a)) -> a -> ST s (STLoc s a))
  -> (forall s . M.STVector s a -> STLoc s a)
  -> M.STVector s a
  -> STLoc s a
locsM wv fl0' ts | M.null ts = fl0' ts
                 | otherwise = V.head fls ts where
  n = M.length ts
  fls :: V.Vector (M.STVector s a -> STLoc s a)
  fls = V.generate n $ \i ts ->
    let upd fl t = fl ts <$ M.write ts i t
        l = fls ! (i - 1)
        x = fls ! i
        r = fls ! min (i + 1) n
    in  AtS (M.read ts i) (wv (upd x)) (upd fl0') (upd l) (upd r)

locsV
  :: ((a -> Loc a) -> a -> Loc a)
  -> (V.Vector a -> Loc a)
  -> V.Vector a
  -> Loc a
locsV wv fl0' ts | V.null ts = fl0' ts
                 | otherwise = V.head fls ts where
  n   = length ts
  fls = V.generate n $ \i ts ->
    let upd fl t = fl (V.take i ts <> V.singleton t <> V.drop (i + 1) ts)
        l = fls ! (i - 1)
        x = fls ! i
        r = fls ! min (i + 1) n
    in  At (ts ! i) (wv (upd x)) (upd fl0') (upd l) (upd r)

locs
  :: forall a . ((a -> Loc a) -> a -> Loc a) -> ([a] -> Loc a) -> [a] -> Loc a
locs wv fl0' [] = fl0' []
locs wv fl0' ts = NE.head fls ts where
  n :: Int
  n = length ts
  fls :: NE.NonEmpty ([a] -> Loc a)
  fls = NE.map
    (\i ts -> At (ts !! i)
                 (wv (upd i (fls NE.!! i)))
                 (upd i fl0')
                 (upd i (fls NE.!! (i - 1)))
                 (upd i (fls NE.!! min (i + 1) n)))
    (NE.fromList [0 ..])
   where
    upd :: Int -> ([a] -> Loc a) -> a -> Loc a
    upd i fl t = fl (take i ts ++ t : drop (i + 1) ts)

explore :: (a -> Weaver a) -> a -> Loc a
explore wv = fr where fr t = At t (call wv fr) fr fr fr

data STLoc s a = AtS
  { itS     :: ST s a
  , fdownS  :: a -> ST s (STLoc s a)
  , fupS    :: a -> ST s (STLoc s a)
  , fleftS  :: a -> ST s (STLoc s a)
  , frightS :: a -> ST s (STLoc s a)
  }

data Loc a = At
  { it     :: a
  , fdown  :: a -> Loc a
  , fup    :: a -> Loc a
  , fleft  :: a -> Loc a
  , fright :: a -> Loc a
  }

down, up, left, right :: Loc a -> Loc a
down  = fdown  <*> it
up    = fup    <*> it
left  = fleft  <*> it
right = fright <*> it

data Term
  = Var String
  | Abs String Term
  | App Term Term
  | If Term Term Term
  deriving Show

-- con0 :: (a -> Weaver a) -> a                                 -> Weaver a
-- con1 :: (a -> Weaver a) -> (a -> a)           -> a           -> Weaver a
-- con2 :: (a -> Weaver a) -> (a -> a -> a)      -> a -> a      -> Weaver a
-- con3 :: (a -> Weaver a) -> (a -> a -> a -> a) -> a -> a -> a -> Weaver a
-- con0 wv k          = W (\fl0 -> loc0 (call wv) (             fl0 k           ))
-- con1 wv k t1       = W (\fl0 -> loc1 (call wv) (\t1 ->       fl0 (k t1)      ) t1)
-- con2 wv k t1 t2    = W (\fl0 -> loc2 (call wv) (\t1 t2 ->    fl0 (k t1 t2)   ) t1 t2)
-- con3 wv k t1 t2 t3 = W (\fl0 -> loc3 (call wv) (\t1 t2 t3 -> fl0 (k t1 t2 t3)) t1 t2 t3)
-- 
-- loc0 :: ((a -> Loc a) -> a -> Loc a) -> Loc a                                 -> Loc a
-- loc1 :: ((a -> Loc a) -> a -> Loc a) -> (a           -> Loc a) -> a           -> Loc a
-- loc2 :: ((a -> Loc a) -> a -> Loc a) -> (a -> a      -> Loc a) -> a -> a      -> Loc a
-- loc3 :: ((a -> Loc a) -> a -> Loc a) -> (a -> a -> a -> Loc a) -> a -> a -> a -> Loc a
-- loc0 wv fl0' = fl0'
-- loc1 wv fl0' = fl1 where
--   fl1 t1 = At t1 (wv (upd fl1)) (upd fl0') (upd fl1) (upd fl1) where
--     upd fl t1' = fl t1'
-- loc2 wv fl0' = fl1 where
--   fl1 t1 t2 = At t1 (wv (upd fl1)) (upd fl0') (upd fl1) (upd fl2) where
--     upd fl t1' = fl t1' t2
--   fl2 t1 t2 = At t2 (wv (upd fl2)) (upd fl0') (upd fl1) (upd fl2) where
--     upd fl t2' = fl t1 t2'
-- loc3 wv fl0' = fl1 where
--   fl1 t1 t2 t3 = At t1 (wv (upd fl1)) (upd fl0') (upd fl1) (upd fl2) where
--     upd fl t1' = fl t1' t2 t3
--   fl2 t1 t2 t3 = At t2 (wv (upd fl2)) (upd fl0') (upd fl1) (upd fl3) where
--     upd fl t2' = fl t1 t2' t3
--   fl3 t1 t2 t3 = At t3 (wv (upd fl3)) (upd fl0') (upd fl2) (upd fl3) where
--     upd fl t3' = fl t1 t2 t3'
-- 
-- weave :: Term -> Weaver Term
-- weave (Var s)          = con0 weave (Var s)
-- weave (Abs s t1)       = con1 weave (Abs s) t1
-- weave (App   t1 t2)    = con2 weave App     t1 t2
-- weave (If    t1 t2 t3) = con3 weave If      t1 t2 t3

weave :: Term -> Weaver Term
weave (Var s   ) = consV weave (\ts -> Var s) (V.fromList [])
weave (Abs s t1) = consV weave (\ts -> Abs s (ts ! 0)) (V.fromList [t1])
weave (App t1 t2) =
  consV weave (\ts -> App (ts ! 0) (ts ! 1)) (V.fromList [t1, t2])
weave (If t1 t2 t3) =
  consV weave (\ts -> If (ts ! 0) (ts ! 1) (ts ! 2)) (V.fromList [t1, t2, t3])

ex1 :: Term
ex1 = Abs
  "n"
  (If
    (App (App (Var "=") (Var "n")) (Var "0"))
    (Var "1")
    (App (App (Var "+") (Var "n"))
         (App (Var "fac") (App (Var "pred") (Var "n")))
    )
  )

ex1' :: Term
ex1' = explore weave ex1
  & down
  & down
  & right
  & right
  & down
  & down
  & flip fup (Var "*")
  & up
  & up
  & up
  & up
  & it

-- top :: Term -> Loc Term
-- top = fr where fr t = At t (weave fr) fr fr fr
-- 
-- weave :: (Term -> Loc Term) -> Term -> Loc Term
-- weave fl0 = \case
--   Var s        -> loc0 weave (                fl0 (Var s            ))
--   Abs s t1     -> loc1 weave (\t1'         -> fl0 (Abs s t1'        )) t1
--   App t1 t2    -> loc2 weave (\t1' t2'     -> fl0 (App   t1' t2'    )) t1 t2
--   If  t1 t2 t3 -> loc3 weave (\t1' t2' t3' -> fl0 (If    t1' t2' t3')) t1 t2 t3

-- repeatFirst :: [a] -> [a]
-- repeatFirst [] = []
-- repeatFirst (x:xs) = x:x:xs
-- 
-- repeatLast :: [a] -> [a]
-- repeatLast xs = foldr (\x go _ -> x : go [x]) id xs []
-- 
-- locs :: (Loc a -> a -> Loc a) -> Loc a -> [a] -> Loc a
-- locs wv l0 [] = l0
-- locs wv l0 ts = head ls where
--   ls = zipWith4 (\lft t l rgt -> At t (wv l t) l0 lft rgt)
--                 (repeatFirst ls)
--                 ts
--                 ls
--                 (repeatLast (tail ls))