Experiment, Optimal reduction with justification (and sausages)

Linc.hs

module Linc where

import Data.STRef
import Control.Monad.ST
import Control.Monad ((>=>))

data Structure = Structure
    { term       :: Tm Int Int
    , link_count :: Int
    , opt_count  :: Int
    , cuts       :: Cuts (Tm Int Int) }
    deriving (Show, Eq)

norm :: Structure -> Structure
norm struct = runST $ do
    (term, cuts) <- struct_eval struct
    struct_cut_full cuts
    struct_reify (term, [])

norm_nce :: Int -> Structure -> Structure
norm_nce n struct = runST $ do
    (term, cuts) <- struct_eval struct
    cuts' <- nest n struct_cut_once cuts
    struct_reify (term, cuts')

struct_eval :: Structure -> ST s (STerm s, Cuts (STerm s))
struct_eval struct = do
    link_env <- fresh (link_count struct)
    opt_env <- fresh_opt (opt_count struct)
    let ev = eval link_env opt_env
    let prep (i,j) = (ev i, ev j)
    pure (ev (term struct), fmap prep (cuts struct))

struct_cut_once :: Cuts (STerm s) -> ST s (Cuts (STerm s))
struct_cut_once [] = pure []
struct_cut_once ((x,y):oldcuts) = do
    cuts <- cut x y
    pure (cuts ++ oldcuts)

struct_cut_full :: Cuts (STerm s) -> ST s ()
struct_cut_full [] = pure ()
struct_cut_full cuts = struct_cut_once cuts >>= struct_cut_full

struct_reify :: (STerm s, Cuts (STerm s)) -> ST s Structure
struct_reify (sterm, cuts) = do
    let renv = ([], [])
    (tm, renv') <- readback sterm renv
    (cs, renv'') <- readback_cuts cuts renv'
    pure (Structure
        { term = tm
        , link_count = length (fst renv'')
        , opt_count = length (snd renv'')
        , cuts = cs })

readback_cuts :: Cuts (STerm s) -> REnv s -> ST s (Cuts (Tm Int Int), REnv s)
readback_cuts [] renv = pure ([], renv)
readback_cuts ((x,y):xs) renv = do
    (x', renv') <- readback x renv
    (y', renv'') <- readback y renv'
    (cs, renv''') <- readback_cuts xs renv''
    pure ((x',y'):cs, renv''')

data J s
    = Decided Bool
    | Duplicated (Maybe (STRef s (J s))) (STRef s (J s)) (STRef s (J s))
    | Erased
    | Instantiated
    | Rescoped (STRef s (J s)) (STRef s (J s))
    | Unbound

data Tm a b
    = Link a
    | Opt b (Tm a b) (Tm a b)
    | On  b (Tm a b) (Tm a b)
    | Exp b (Tm a b)
    | At  b (Tm a b)
    | Fan b (Tm a b) (Tm a b)
    | Pair  (Tm a b) (Tm a b)
    | Index Bool (Tm a b)
    | Read       (Tm a b)
    | Drop
    deriving (Show, Eq)

type STerm s = Tm (Box s) (STRef s (J s))
newtype Box s = Box (STRef s (Maybe (STerm s)))

eval :: [Box s] -> [STRef s (J s)] -> Tm Int Int -> Tm (Box s) (STRef s (J s))
eval env opt (Link i) = Link (env !! i)
eval env opt (Opt j x y) = Opt (opt !! j) (eval env opt x) (eval env opt y)
eval env opt (On  j x y) = On  (opt !! j) (eval env opt x) (eval env opt y)
eval env opt (Exp j x)   = Exp (opt !! j) (eval env opt x)
eval env opt (At  j x)   = At  (opt !! j) (eval env opt x)
eval env opt (Fan j x y) = Fan (opt !! j) (eval env opt x) (eval env opt y)
eval env opt (Pair x y)  = Pair           (eval env opt x) (eval env opt y)
eval env opt (Index b x) = Index b (eval env opt x)
eval env opt (Read x)    = Read    (eval env opt x)
eval env opt Drop        = Drop

blank :: ST s (Box s)
blank = Box `fmap` newSTRef Nothing

unbound :: ST s (STRef s (J s))
unbound = newSTRef Unbound
 
fresh :: Int -> ST s [Box s]
fresh n = sequence (replicate n blank)

fresh_opt :: Int -> ST s [STRef s (J s)]
fresh_opt n = sequence (replicate n unbound)

type Cuts a = [(a, a)]

cut :: STerm s -> STerm s -> ST s (Cuts (STerm s))
cut a b = do
    (x,c1) <- refine a
    (y,c2) <- refine b
    if priority x <= priority y
    then do c3 <- slice x y
            pure (c1++c2++c3)
    else do c3 <- slice y x
            pure (c1++c2++c3)

priority :: Tm a b -> Int
priority (Link _)    = 0
priority Drop        = 1
priority (Fan _ _ _) = 1
priority (Index _ _) = 2
priority (Read    _) = 2
priority (Pair  _ _) = 2
priority (Opt _ _ _) = 4
priority (On  _ _ _) = 4
priority (Exp _ _)   = 4
priority (At  _ _)   = 5

cut_refine :: STerm s -> Cuts (STerm s)
         -> ST s (Tm (Box s) (STRef s (J s)), Cuts (STerm s))
cut_refine x cuts = do
    (y,cuts') <- refine x
    pure (y, cuts' ++ cuts)

refine :: STerm s -> ST s (Tm (Box s) (STRef s (J s)), Cuts (STerm s))
refine (Fan i x y) = do
    Duplicated _ a b <- readSTRef i
    b1 <- readSTRef a >>= erased
    b2 <- readSTRef b >>= erased
    case (b1,b2) of
        (False,False) -> pure (Fan i x y, [])
        (True, False) -> cut_refine y [(Drop,x)]
        (False, True) -> cut_refine x [(Drop,y)]
        (True,  True) -> pure (Drop, [(Drop,x),(Drop,y)])
refine (Opt i x y) = do
    status <- readSTRef i
    case status of
        Duplicated k a b -> do
            c0 <- fmap Link blank
            c1 <- fmap Link blank
            c2 <- fmap Link blank
            c3 <- fmap Link blank
            cut_refine ((maybe Pair Fan k) (Opt a c0 c1) (Opt b c2 c3))
                       [(Fan i c0 c2, x), (Fan i c1 c3, y)]
        Erased  -> pure (Drop, [(Drop,x)])
        Unbound -> pure (Opt i x y, [])
refine (On  i x y) = do
    status <- readSTRef i
    case status of
        Decided False -> cut_refine x [(Drop,y)]
        Decided True  -> cut_refine y [(Drop,x)]
        Duplicated k a b -> do
            c0 <- fmap Link blank
            c1 <- fmap Link blank
            c2 <- fmap Link blank
            c3 <- fmap Link blank
            cut_refine ((maybe Pair Fan k) (On a c0 c1) (On b c2 c3))
                       [(Fan i c0 c2, x), (Fan i c1 c3, y)]
        Erased  -> pure (Drop, [(Drop,x)])
        Unbound -> pure (On i x y, [])
refine (Exp i x) = do
    status <- readSTRef i
    case status of
        Duplicated k a b -> do
            c0 <- fmap Link blank
            c2 <- fmap Link blank
            cut_refine ((maybe Pair Fan k) (Exp a c0) (Exp b c2))
                       [(Fan i c0 c2, x)]
        Unbound -> pure (Exp i x, [])
refine (At  i x) = do
    status <- readSTRef i
    case status of
        Duplicated k a b -> do
            c0 <- fmap Link blank
            c2 <- fmap Link blank
            cut_refine ((maybe Pair Fan k) (At a c0) (At b c2))
                       [(Fan i c0 c2, x)]
        Erased       -> pure (Drop, [(Drop,x)])
        Instantiated -> refine x
        Rescoped j k -> refine (At j (At k x))
        Unbound      -> pure (Exp i x, [])
refine others = pure (others, [])

slice :: STerm s -> STerm s -> ST s (Cuts (STerm s))
slice (Link (Box ref)) y = do
    state <- readSTRef ref
    case state of
        Nothing -> writeSTRef ref (Just y) >> pure []
        Just x  -> pure [(x,y)]
slice Drop Drop        = pure []
slice Drop (Fan _ x y) = pure [(Drop,x), (Drop,y)]
slice Drop (Index _ x) = pure [(Drop,x)]
slice Drop (Read    x) = pure [(Drop,x)]
slice Drop (Pair  x y) = pure [(Drop,x), (Drop,y)]
slice Drop (Opt i x y) = writeSTRef i Erased >> pure [(Drop,x), (Drop,y)]
slice Drop (On  _ x y) = pure [(Drop,x), (Drop,y)]
slice Drop (Exp i   x) = writeSTRef i Erased >> pure [(Drop,x)]
slice Drop (At i   x) = pure [(Drop,x)]
slice (Fan i x y) (Fan j z w) | (i == j) = do (++) <$> slice x z <*> slice y w
slice (Fan i x y) (Fan j z w) = commute2x2 (Fan i) x y (Fan j) z w
slice (Fan i x y) (Index b z) = commute2x1 (Fan i) x y (Index b) z
slice (Fan i x y) (Read    z) = commute2x1 (Fan i) x y Read z
slice (Fan i x y) (Pair  z w) = commute2x2 (Fan i) x y Pair z w
slice (Fan i x y) (Opt j z w) = do
    (p,q) <- duplicate (Just i) j
    commute2x2D (Fan j) x y (Opt p) (Opt q) z w
slice (Fan i x y) (On j z w) = do
    (p,q) <- duplicate (Just i) j
    commute2x2D (Fan j) x y (On p) (On q) z w
slice (Fan i x y) (Exp j z) = do
    (p,q) <- duplicate (Just i) j
    commute2x1D (Fan j) x y (Exp p) (Exp q) z
slice (Fan i x y) (At j z) = do
    (p,q) <- duplicate (Just i) j
    commute2x1D (Fan j) x y (At p) (At q) z
slice (Index False z) (Opt i x y) = writeSTRef i (Decided False) >> pure [(z,x)]
slice (Index True  z) (Opt i x y) = writeSTRef i (Decided True) >> pure [(z,y)]
slice (Read z)        (Exp i x) = writeSTRef i Instantiated >> pure [(z,x)]
slice (Pair x y)  (Pair z w) = pure [(x,z), (y,w)]
slice (Pair x y)  (On i z w) = commute2x2 Pair x y (On i) z w
slice (Pair x y)  (Exp j z) = do
    (p,q) <- duplicate Nothing j
    commute2x1D (Fan j) x y (Exp p) (Exp q) z
slice (Exp i x) (At j z) = do
    a <- newSTRef Unbound
    writeSTRef i (Rescoped j a)
    pure [(Exp a x, z)]
slice x y = error (show (info x, info y))

info :: STerm s -> String
info (Link _)    = "Link"
info (Drop)      = "Drop"
info (Fan _ _ _) = "Fan"
info (Index _ _) = "Index"
info (Read    _) = "Read"
info (Pair  _ _) = "Pair"
info (Opt _ _ _) = "Opt"
info (On  _ _ _) = "On"
info (Exp _ _)   = "Exp"
info (At  _ _)   = "At"

commute2x2D f x y g h z w = do
    a <- fmap Link blank
    b <- fmap Link blank
    c <- fmap Link blank
    d <- fmap Link blank
    pure [(x, f a b), (y, f c d), (g a c, z), (h b d, w)]

commute2x2 f x y g z w = commute2x2D f x y g g z w

commute2x1D f x y g h z = do
    a <- fmap Link blank
    b <- fmap Link blank
    pure [(z, f a b), (g a, x), (h b, y)]

commute2x1 f x y g z = commute2x1D f x y g g z

commute1x1 f x g z = do
    a <- fmap Link blank
    pure [(z, f a), (g a, x)]

duplicate :: Maybe (STRef s (J s)) -> STRef s (J s)
          -> ST s (STRef s (J s), STRef s (J s))
duplicate k i = do
    a <- newSTRef Unbound
    b <- newSTRef Unbound
    writeSTRef i (Duplicated k a b)
    pure (a,b)

erased :: J s -> ST s Bool
erased (Decided _)      = pure False
erased (Duplicated _ x y) = (&&) <$> (readSTRef x >>= erased)
                                 <*> (readSTRef y >>= erased)
erased Erased           = pure True
erased Instantiated     = pure False
erased Unbound          = pure False
 
type REnv s = ([STRef s (Maybe (STerm s))], [STRef s (J s)])

readback_done :: (Tm Int Int, REnv s) -> (Tm Int Int, Int, Int)
readback_done (res, (e1,e2)) = (res, length e1, length e2)

readback :: STerm s -> REnv s -> ST s (Tm Int Int, REnv s)
readback structure env = refine structure >>= \(s,_) -> readback' s env

readback' :: STerm s -> REnv s -> ST s (Tm Int Int, REnv s)
readback' (Link (Box ref)) env = do
    state <- readSTRef ref
    case state of
        Nothing -> let (i, env') = plot0 ref env
                   in pure (Link i, env')
        Just y  -> readback y env
readback' (Opt k x y) env = let (i, env') = plot1 k env
    in do (x', env'') <- readback x env'
          (y', env''') <- readback y env''
          pure (Opt i x' y', env''')
readback' (On  k x y) env = let (i, env') = plot1 k env
    in do (x', env'') <- readback x env'
          (y', env''') <- readback y env''
          pure (On  i x' y', env''')
readback' (Exp k x) env = let (i, env') = plot1 k env
    in do (x', env'') <- readback x env'
          pure (Exp i x', env'')
readback' (At  k x) env = let (i, env') = plot1 k env
    in do (x', env'') <- readback x env'
          pure (At  i x', env'')
readback' (Fan k x y) env = let (i, env') = plot1 k env
    in do (x', env'') <- readback x env'
          (y', env''') <- readback y env''
          pure (Fan i x' y', env''')
readback' (Pair x y) env = do
    (x', env') <- readback x env
    (y', env'') <- readback y env'
    pure (Pair x' y', env'')
readback' (Index b x) env = do
    (x', env') <- readback x env
    pure (Index b x', env')
readback' (Read x) env = do
    (x', env') <- readback x env
    pure (Read x', env')
readback' Drop env = pure (Drop, env)

plot0 :: Eq t1 => t1 -> ([t1], t) -> (Int, ([t1], t))
plot0 ref (e1,e2) = let (k,e1') = plot ref e1 in (k, (e1',e2))

plot1 :: Eq t1 => t1 -> (t, [t1]) -> (Int, (t, [t1]))
plot1 ref (e1,e2) = let (k,e2') = plot ref e2 in (k, (e1,e2'))

plot :: (Eq a) => a -> [a] -> (Int, [a])
plot y [] = (0,[y])
plot y (x:xs) | (x==y) = (0,x:xs)
plot y (x:xs)          = let (k,zs) = plot y xs in (k+1, x:zs)

nest :: Monad f => Int -> (a -> f a) -> a -> f a
nest n f = foldr (>=>) pure (replicate n f)