{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}

module Unison.ABT.Normalized
  ( ABT (..),
    Term (.., TAbs, TTm, TAbss),
    Align (..),
    alpha,
    freshen,
    Renaming (..),
    isEmptyRenaming,
    freshenBinder,
    freshenBinders,
    pruneRenaming,
    renameVar,
    mapping,
    avoiding,
    mappingAndAvoiding,
    renames,
    renamesAvoiding,
    renamesAndFreshen0,
    rename,
    transform,
    visit,
    visitPure,
  )
where

import Data.Bifoldable
import Data.Bifunctor
import Data.Functor.Identity (Identity (..))
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe)
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Traversable (mapAccumL)
import Unison.ABT (Var (..))

-- ABTs with support for 'normalized' structure where only variables
-- may occur at some positions. This is accomplished by passing the
-- variable type to the base functor.
data ABT f v
  = Abs v (Term f v)
  | Tm (f v (Term f v))

data Term f v = Term
  { forall (f :: * -> * -> *) v. Term f v -> Set v
freeVars :: Set v,
    forall (f :: * -> * -> *) v. Term f v -> ABT f v
out :: ABT f v
  }

instance
  (forall a b. (Show a) => (Show b) => Show (f a b), Show v) =>
  Show (ABT f v)
  where
  showsPrec :: Int -> ABT f v -> ShowS
showsPrec Int
p ABT f v
a = Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
9) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ case ABT f v
a of
    Abs v
v Term f v
tm ->
      String -> ShowS
showString String
"Abs "
        ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> v -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
10 v
v
        ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
" "
        ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Term f v -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
10 Term f v
tm
    Tm f v (Term f v)
e -> String -> ShowS
showString String
"Tm " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> f v (Term f v) -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
10 f v (Term f v)
e

instance
  (forall a b. (Show a) => (Show b) => Show (f a b), Show v) =>
  Show (Term f v)
  where
  showsPrec :: Int -> Term f v -> ShowS
showsPrec Int
p (Term Set v
_ ABT f v
e) =
    Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
9) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ String -> ShowS
showString String
"Term " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ABT f v -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
10 ABT f v
e

instance
  (forall a b. (Eq a) => (Eq b) => Eq (f a b), Bifunctor f, Bifoldable f, Var v) =>
  Eq (ABT f v)
  where
  Abs v
v1 Term f v
e1 == :: ABT f v -> ABT f v -> Bool
== Abs v
v2 Term f v
e2
    | v
v1 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v2 = Term f v
e1 Term f v -> Term f v -> Bool
forall a. Eq a => a -> a -> Bool
== Term f v
e2
    | Bool
otherwise = Term f v
e1 Term f v -> Term f v -> Bool
forall a. Eq a => a -> a -> Bool
== v -> v -> Term f v -> Term f v
forall v (f :: * -> * -> *).
(Var v, Ord v, Bifunctor f, Bifoldable f) =>
v -> v -> Term f v -> Term f v
rename v
v2 v
v1 Term f v
e2
  Tm f v (Term f v)
e1 == Tm f v (Term f v)
e2 = f v (Term f v)
e1 f v (Term f v) -> f v (Term f v) -> Bool
forall a. Eq a => a -> a -> Bool
== f v (Term f v)
e2
  ABT f v
_ == ABT f v
_ = Bool
False

instance
  (forall a b. (Eq a) => (Eq b) => Eq (f a b), Bifunctor f, Bifoldable f, Var v) =>
  Eq (Term f v)
  where
  Term Set v
_ ABT f v
abt1 == :: Term f v -> Term f v -> Bool
== Term Set v
_ ABT f v
abt2 = ABT f v
abt1 ABT f v -> ABT f v -> Bool
forall a. Eq a => a -> a -> Bool
== ABT f v
abt2

pattern TAbs :: (Var v) => v -> Term f v -> Term f v
pattern $mTAbs :: forall {r} {v} {f :: * -> * -> *}.
Var v =>
Term f v -> (v -> Term f v -> r) -> ((# #) -> r) -> r
$bTAbs :: forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
TAbs u bd <-
  Term _ (Abs u bd)
  where
    TAbs v
u Term f v
bd = Set v -> ABT f v -> Term f v
forall (f :: * -> * -> *) v. Set v -> ABT f v -> Term f v
Term (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.delete v
u (Term f v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
freeVars Term f v
bd)) (v -> Term f v -> ABT f v
forall (f :: * -> * -> *) v. v -> Term f v -> ABT f v
Abs v
u Term f v
bd)

pattern TTm :: (Var v, Bifoldable f) => f v (Term f v) -> Term f v
pattern $mTTm :: forall {r} {v} {f :: * -> * -> *}.
(Var v, Bifoldable f) =>
Term f v -> (f v (Term f v) -> r) -> ((# #) -> r) -> r
$bTTm :: forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
TTm bd <-
  Term _ (Tm bd)
  where
    TTm f v (Term f v)
bd = Set v -> ABT f v -> Term f v
forall (f :: * -> * -> *) v. Set v -> ABT f v -> Term f v
Term ((v -> Set v) -> (Term f v -> Set v) -> f v (Term f v) -> Set v
forall m a b. Monoid m => (a -> m) -> (b -> m) -> f a b -> m
forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap v -> Set v
forall a. a -> Set a
Set.singleton Term f v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
freeVars f v (Term f v)
bd) (f v (Term f v) -> ABT f v
forall (f :: * -> * -> *) v. f v (Term f v) -> ABT f v
Tm f v (Term f v)
bd)

{-# COMPLETE TAbs, TTm #-}

class (Bifoldable f, Bifunctor f) => Align f where
  align ::
    (Applicative g) =>
    (vl -> vr -> g vs) ->
    (el -> er -> g es) ->
    f vl el ->
    f vr er ->
    Maybe (g (f vs es))

alphaErr ::
  (Align f) => (Var v) => Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) a
alphaErr :: forall (f :: * -> * -> *) v a.
(Align f, Var v) =>
Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) a
alphaErr Map v v
un Term f v
tml Term f v
tmr = (Term f v, Term f v) -> Either (Term f v, Term f v) a
forall a b. a -> Either a b
Left (Term f v
tml, Renaming v -> Term f v -> Term f v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
renamesAndFreshen0 Renaming v
rn Term f v
tmr)
  where
    rn :: Renaming v
rn = Map v v -> Renaming v
forall v. Var v => Map v v -> Renaming v
mapping Map v v
un

-- Checks if two terms are equal up to a given variable renaming. The
-- renaming should map variables in the right hand term to the
-- equivalent variable in the left hand term.
alpha :: (Align f) => (Var v) => Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) ()
alpha :: forall (f :: * -> * -> *) v.
(Align f, Var v) =>
Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) ()
alpha Map v v
un (TAbs v
u Term f v
tml) (TAbs v
v Term f v
tmr) =
  Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) ()
forall (f :: * -> * -> *) v.
(Align f, Var v) =>
Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) ()
alpha (v -> v -> Map v v -> Map v v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert v
v v
u ((v -> Bool) -> Map v v -> Map v v
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
u) Map v v
un)) Term f v
tml Term f v
tmr
alpha Map v v
un tml :: Term f v
tml@(TTm f v (Term f v)
bdl) tmr :: Term f v
tmr@(TTm f v (Term f v)
bdr)
  | Just Either (Term f v, Term f v) (f () ())
sub <- (v -> v -> Either (Term f v, Term f v) ())
-> (Term f v -> Term f v -> Either (Term f v, Term f v) ())
-> f v (Term f v)
-> f v (Term f v)
-> Maybe (Either (Term f v, Term f v) (f () ()))
forall (g :: * -> *) vl vr vs el er es.
Applicative g =>
(vl -> vr -> g vs)
-> (el -> er -> g es) -> f vl el -> f vr er -> Maybe (g (f vs es))
forall (f :: * -> * -> *) (g :: * -> *) vl vr vs el er es.
(Align f, Applicative g) =>
(vl -> vr -> g vs)
-> (el -> er -> g es) -> f vl el -> f vr er -> Maybe (g (f vs es))
align v -> v -> Either (Term f v, Term f v) ()
av (Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) ()
forall (f :: * -> * -> *) v.
(Align f, Var v) =>
Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) ()
alpha Map v v
un) f v (Term f v)
bdl f v (Term f v)
bdr = () ()
-> Either (Term f v, Term f v) (f () ())
-> Either (Term f v, Term f v) ()
forall a b.
a -> Either (Term f v, Term f v) b -> Either (Term f v, Term f v) a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Either (Term f v, Term f v) (f () ())
sub
  where
    av :: v -> v -> Either (Term f v, Term f v) ()
av v
u v
v
      | Bool -> (v -> Bool) -> Maybe v -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
u) (v -> Map v v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup v
v Map v v
un) = () -> Either (Term f v, Term f v) ()
forall a. a -> Either (Term f v, Term f v) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) ()
forall (f :: * -> * -> *) v a.
(Align f, Var v) =>
Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) a
alphaErr Map v v
un Term f v
tml Term f v
tmr
alpha Map v v
un Term f v
tml Term f v
tmr = Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) ()
forall (f :: * -> * -> *) v a.
(Align f, Var v) =>
Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) a
alphaErr Map v v
un Term f v
tml Term f v
tmr

unabss :: (Var v) => Term f v -> ([v], Term f v)
unabss :: forall v (f :: * -> * -> *). Var v => Term f v -> ([v], Term f v)
unabss (TAbs v
v (Term f v -> ([v], Term f v)
forall v (f :: * -> * -> *). Var v => Term f v -> ([v], Term f v)
unabss -> ([v]
vs, Term f v
bd))) = (v
v v -> [v] -> [v]
forall a. a -> [a] -> [a]
: [v]
vs, Term f v
bd)
unabss Term f v
bd = ([], Term f v
bd)

pattern TAbss :: (Var v) => [v] -> Term f v -> Term f v
pattern $mTAbss :: forall {r} {v} {f :: * -> * -> *}.
Var v =>
Term f v -> ([v] -> Term f v -> r) -> ((# #) -> r) -> r
$bTAbss :: forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
TAbss vs bd <-
  (unabss -> (vs, bd))
  where
    TAbss [v]
vs Term f v
bd = (v -> Term f v -> Term f v) -> Term f v -> [v] -> Term f v
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr v -> Term f v -> Term f v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
TAbs Term f v
bd [v]
vs

{-# COMPLETE TAbss #-}

-- Renaming data.
--
-- `renames` contains an actual mapping from old variables to new
-- variables, which is used for the actual substitution.
--
-- `conflicts` stores information about which variables should be
-- avoided. The value at a variable `u` should _at least_ be the
-- number of variables in `renames` that are being substituted _to_
-- `u`, because substituting under a binder for `u` will capture those
-- substitutions. However, `conflicts` can be bootstrapped with extra
-- counts to avoid ambient variables as well, so that it is possible
-- to substitute/rewrite expressions without introducing variable
-- captures.
data Renaming v = RN
  { forall v. Renaming v -> Map v Int
conflicts :: Map v Int,
    forall v. Renaming v -> Map v v
renamings :: Map v v
  }

-- Creates a Renaming that will avoid the given set of variables.
avoiding :: Set v -> Renaming v
avoiding :: forall v. Set v -> Renaming v
avoiding Set v
avoid = Map v Int -> Map v v -> Renaming v
forall v. Map v Int -> Map v v -> Renaming v
RN ((v -> Int) -> Set v -> Map v Int
forall k a. (k -> a) -> Set k -> Map k a
Map.fromSet (Int -> v -> Int
forall a b. a -> b -> a
const Int
1) Set v
avoid) Map v v
forall k a. Map k a
Map.empty

mapping :: (Var v) => Map v v -> Renaming v
mapping :: forall v. Var v => Map v v -> Renaming v
mapping Map v v
rn = Map v Int -> Map v v -> Renaming v
forall v. Map v Int -> Map v v -> Renaming v
RN Map v Int
cf Map v v
rn
  where
    cf :: Map v Int
cf = (Int -> Int -> Int) -> [(v, Int)] -> Map v Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([(v, Int)] -> Map v Int)
-> ([v] -> [(v, Int)]) -> [v] -> Map v Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (v -> (v, Int)) -> [v] -> [(v, Int)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (,Int
1) ([v] -> Map v Int) -> [v] -> Map v Int
forall a b. (a -> b) -> a -> b
$ Map v v -> [v]
forall k a. Map k a -> [a]
Map.elems Map v v
rn

mappingAndAvoiding :: (Var v) => Map v v -> Set v -> Renaming v
mappingAndAvoiding :: forall v. Var v => Map v v -> Set v -> Renaming v
mappingAndAvoiding Map v v
rn Set v
avoid = Map v Int -> Map v v -> Renaming v
forall v. Map v Int -> Map v v -> Renaming v
RN Map v Int
cf Map v v
rn
  where
    cf :: Map v Int
cf =
      (Int -> Int -> Int) -> Map v Int -> Map v Int -> Map v Int
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith
        Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)
        ((v -> Int) -> Set v -> Map v Int
forall k a. (k -> a) -> Set k -> Map k a
Map.fromSet (Int -> v -> Int
forall a b. a -> b -> a
const Int
1) Set v
avoid)
        ((Int -> Int -> Int) -> [(v, Int)] -> Map v Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([(v, Int)] -> Map v Int)
-> ([v] -> [(v, Int)]) -> [v] -> Map v Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (v -> (v, Int)) -> [v] -> [(v, Int)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (,Int
1) ([v] -> Map v Int) -> [v] -> Map v Int
forall a b. (a -> b) -> a -> b
$ Map v v -> [v]
forall k a. Map k a -> [a]
Map.elems Map v v
rn)

-- Adjusts a renaming with respect to a remaining set of free
-- variables. Unnecessary renamings are discarded.
pruneRenaming :: (Var v) => Set v -> Renaming v -> Renaming v
pruneRenaming :: forall v. Var v => Set v -> Renaming v -> Renaming v
pruneRenaming Set v
fvs (RN Map v Int
cf Map v v
rn) =
  RN
    { $sel:renamings:RN :: Map v v
renamings = Map v v -> Set v -> Map v v
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.restrictKeys Map v v
rn Set v
fvs,
      $sel:conflicts:RN :: Map v Int
conflicts = (Map v Int -> v -> Map v Int) -> Map v Int -> Map v v -> Map v Int
forall a b k. (a -> b -> a) -> a -> Map k b -> a
Map.foldl' Map v Int -> v -> Map v Int
forall {k} {a}. (Ord k, Ord a, Num a) => Map k a -> k -> Map k a
decrement Map v Int
cf (Map v v -> Map v Int) -> Map v v -> Map v Int
forall a b. (a -> b) -> a -> b
$ Map v v -> Set v -> Map v v
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.withoutKeys Map v v
rn Set v
fvs
    }
  where
    decrement :: Map k a -> k -> Map k a
decrement Map k a
sv k
v = (a -> Maybe a) -> k -> Map k a -> Map k a
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update a -> Maybe a
forall {a}. (Ord a, Num a) => a -> Maybe a
drop k
v Map k a
sv
    drop :: a -> Maybe a
drop a
n
      | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
1 = Maybe a
forall a. Maybe a
Nothing
      | Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
1)

renameVar :: (Var v) => Renaming v -> v -> v
renameVar :: forall v. Var v => Renaming v -> v -> v
renameVar (RN Map v Int
_ Map v v
rn) v
u = v -> v -> Map v v -> v
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault v
u v
u Map v v
rn

-- Tests if the renaming is empty in the sense that it will never
-- cause bound variables to be renamed. This is _not_ just a test of
-- whether the substitutions are empty, because the conflicts can
-- cause variables to need freshening even without variable
-- substitutions.
isEmptyRenaming :: Renaming v -> Bool
isEmptyRenaming :: forall v. Renaming v -> Bool
isEmptyRenaming = Map v Int -> Bool
forall a. Map v a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Map v Int -> Bool)
-> (Renaming v -> Map v Int) -> Renaming v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Renaming v -> Map v Int
forall v. Renaming v -> Map v Int
conflicts

-- Freshens a bound variable with regard to a renaming, yielding the
-- fresh variable and a renaming appropriate for the term within the
-- binder. The `Set` should be the free variables of the expression
-- within the binder, for proper freshening.
freshenBinder :: (Var v) => Set v -> Renaming v -> v -> (Renaming v, v)
freshenBinder :: forall v. Var v => Set v -> Renaming v -> v -> (Renaming v, v)
freshenBinder Set v
fvs rn0 :: Renaming v
rn0@(RN Map v Int
cf Map v v
rn) v
u = (Renaming v
rn', v
u')
  where
    -- if u conflicts with the renaming, freshen it
    u' :: v
u'
      | v
u v -> Map v Int -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`Map.member` Map v Int
cf = Set v -> v -> v
forall v. Var v => Set v -> v -> v
freshIn (Set v
fvs Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Map v Int -> Set v
forall k a. Map k a -> Set k
Map.keysSet Map v Int
cf) v
u
      | Bool
otherwise = v
u

    -- if u needs to be renamed, and it actually occurs in the body,
    -- add it to the Renaming.
    rn' :: Renaming v
rn'
      | v
u v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
u' Bool -> Bool -> Bool
&& v
u v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set v
fvs =
          RN
            { $sel:conflicts:RN :: Map v Int
conflicts = (Int -> Int -> Int) -> v -> Int -> Map v Int -> Map v Int
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) v
u' Int
1 Map v Int
cf,
              $sel:renamings:RN :: Map v v
renamings = (Maybe v -> Maybe v) -> v -> Map v v -> Map v v
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (Maybe v -> Maybe v -> Maybe v
forall a b. a -> b -> a
const (Maybe v -> Maybe v -> Maybe v) -> Maybe v -> Maybe v -> Maybe v
forall a b. (a -> b) -> a -> b
$ v -> Maybe v
forall a. a -> Maybe a
Just v
u') v
u Map v v
rn
            }
      | Bool
otherwise = Renaming v
rn0

freshenBinders ::
  (Var v) => Set v -> Renaming v -> [v] -> (Renaming v, [v])
freshenBinders :: forall v. Var v => Set v -> Renaming v -> [v] -> (Renaming v, [v])
freshenBinders Set v
fvs = (Renaming v -> v -> (Renaming v, v))
-> Renaming v -> [v] -> (Renaming v, [v])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (Set v -> Renaming v -> v -> (Renaming v, v)
forall v. Var v => Set v -> Renaming v -> v -> (Renaming v, v)
freshenBinder Set v
fvs)

-- Simultaneous variable renaming and freshening implementation.
--
-- subvs0 is a count of the number of conflicts associated with a
-- variable. There are two sources of conflicts.
--
--   1. A variable is being renamed _to_ the given variable
--   2. We want to avoid capturing the variable for other reasons
--
-- So, if you initially call `renamesAndFreshen0` with a higher count
-- for `v` then there are variables being renamed to `v`, all bound
-- occurrences of `v` will be freshened regardless of whether any
-- actual renamings are left.
--
-- rnv0 is the variable renaming map.
renamesAndFreshen0 ::
  (Var v, Bifunctor f, Bifoldable f) =>
  Renaming v ->
  Term f v ->
  Term f v
renamesAndFreshen0 :: forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
renamesAndFreshen0 Renaming v
rn0 Term f v
tm = case Term f v
tm of
  TAbs v
u Term f v
body
    | (Renaming v
rn, v
u') <- Set v -> Renaming v -> v -> (Renaming v, v)
forall v. Var v => Set v -> Renaming v -> v -> (Renaming v, v)
freshenBinder (Term f v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
freeVars Term f v
body) Renaming v
rn v
u,
      v
u v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
u' Bool -> Bool -> Bool
|| Bool -> Bool
not (Renaming v -> Bool
forall v. Renaming v -> Bool
isEmptyRenaming Renaming v
rn) ->
        v -> Term f v -> Term f v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
TAbs v
u' (Renaming v -> Term f v -> Term f v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
renamesAndFreshen0 Renaming v
rn Term f v
body)
  TTm f v (Term f v)
body
    | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Renaming v -> Bool
forall v. Renaming v -> Bool
isEmptyRenaming Renaming v
rn ->
        f v (Term f v) -> Term f v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
TTm (f v (Term f v) -> Term f v) -> f v (Term f v) -> Term f v
forall a b. (a -> b) -> a -> b
$ (v -> v)
-> (Term f v -> Term f v) -> f v (Term f v) -> f v (Term f v)
forall a b c d. (a -> b) -> (c -> d) -> f a c -> f b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Renaming v -> v -> v
forall v. Var v => Renaming v -> v -> v
renameVar Renaming v
rn) (Renaming v -> Term f v -> Term f v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
renamesAndFreshen0 Renaming v
rn) f v (Term f v)
body
  Term f v
_ -> Term f v
tm
  where
    fvs :: Set v
fvs = Term f v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
freeVars Term f v
tm

    rn :: Renaming v
rn = Set v -> Renaming v -> Renaming v
forall v. Var v => Set v -> Renaming v -> Renaming v
pruneRenaming Set v
fvs Renaming v
rn0

-- Freshens the bound variables in a term to avoid capturing variables
-- in the set.
freshen ::
  (Var v, Bifunctor f, Bifoldable f) =>
  Set v ->
  Term f v ->
  Term f v
freshen :: forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Set v -> Term f v -> Term f v
freshen Set v
avoid = Renaming v -> Term f v -> Term f v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
renamesAndFreshen0 (Set v -> Renaming v
forall v. Set v -> Renaming v
avoiding Set v
avoid)

-- Renames some variables while also avoiding a given set of variables
-- for any bindings in the term.
renamesAvoiding ::
  (Var v, Bifunctor f, Bifoldable f) =>
  Set v ->
  Map v v ->
  Term f v ->
  Term f v
renamesAvoiding :: forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Set v -> Map v v -> Term f v -> Term f v
renamesAvoiding Set v
avoid Map v v
rnv =
  Renaming v -> Term f v -> Term f v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
renamesAndFreshen0 (Map v v -> Set v -> Renaming v
forall v. Var v => Map v v -> Set v -> Renaming v
mappingAndAvoiding Map v v
rnv Set v
avoid)

-- Simultaneous variable renaming.
renames ::
  (Var v, Bifunctor f, Bifoldable f) =>
  Map v v ->
  Term f v ->
  Term f v
renames :: forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Map v v -> Term f v -> Term f v
renames Map v v
rnv Term f v
tm = Renaming v -> Term f v -> Term f v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
renamesAndFreshen0 (Map v v -> Renaming v
forall v. Var v => Map v v -> Renaming v
mapping Map v v
rnv) Term f v
tm

rename ::
  (Var v, Ord v, Bifunctor f, Bifoldable f) =>
  v ->
  v ->
  Term f v ->
  Term f v
rename :: forall v (f :: * -> * -> *).
(Var v, Ord v, Bifunctor f, Bifoldable f) =>
v -> v -> Term f v -> Term f v
rename v
old v
new = Renaming v -> Term f v -> Term f v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
renamesAndFreshen0 (Map v v -> Renaming v
forall v. Var v => Map v v -> Renaming v
mapping (Map v v -> Renaming v) -> Map v v -> Renaming v
forall a b. (a -> b) -> a -> b
$ v -> v -> Map v v
forall k a. k -> a -> Map k a
Map.singleton v
old v
new)

transform ::
  (Var v, Bifunctor g, Bifoldable f, Bifoldable g) =>
  (forall a b. f a b -> g a b) ->
  Term f v ->
  Term g v
transform :: forall v (g :: * -> * -> *) (f :: * -> * -> *).
(Var v, Bifunctor g, Bifoldable f, Bifoldable g) =>
(forall a b. f a b -> g a b) -> Term f v -> Term g v
transform forall a b. f a b -> g a b
phi (TTm f v (Term f v)
body) = g v (Term g v) -> Term g v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
TTm (g v (Term g v) -> Term g v)
-> (g v (Term f v) -> g v (Term g v)) -> g v (Term f v) -> Term g v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term f v -> Term g v) -> g v (Term f v) -> g v (Term g v)
forall b c a. (b -> c) -> g a b -> g a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((forall a b. f a b -> g a b) -> Term f v -> Term g v
forall v (g :: * -> * -> *) (f :: * -> * -> *).
(Var v, Bifunctor g, Bifoldable f, Bifoldable g) =>
(forall a b. f a b -> g a b) -> Term f v -> Term g v
transform f a b -> g a b
forall a b. f a b -> g a b
phi) (g v (Term f v) -> Term g v) -> g v (Term f v) -> Term g v
forall a b. (a -> b) -> a -> b
$ f v (Term f v) -> g v (Term f v)
forall a b. f a b -> g a b
phi f v (Term f v)
body
transform forall a b. f a b -> g a b
phi (TAbs v
u Term f v
body) = v -> Term g v -> Term g v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
TAbs v
u (Term g v -> Term g v) -> Term g v -> Term g v
forall a b. (a -> b) -> a -> b
$ (forall a b. f a b -> g a b) -> Term f v -> Term g v
forall v (g :: * -> * -> *) (f :: * -> * -> *).
(Var v, Bifunctor g, Bifoldable f, Bifoldable g) =>
(forall a b. f a b -> g a b) -> Term f v -> Term g v
transform f a b -> g a b
forall a b. f a b -> g a b
phi Term f v
body

visit ::
  (Applicative g, Bifoldable f, Traversable (f v), Var v) =>
  (Term f v -> Maybe (g (Term f v))) ->
  Term f v ->
  g (Term f v)
visit :: forall (g :: * -> *) (f :: * -> * -> *) v.
(Applicative g, Bifoldable f, Traversable (f v), Var v) =>
(Term f v -> Maybe (g (Term f v))) -> Term f v -> g (Term f v)
visit Term f v -> Maybe (g (Term f v))
h Term f v
t = (g (Term f v) -> Maybe (g (Term f v)) -> g (Term f v))
-> Maybe (g (Term f v)) -> g (Term f v) -> g (Term f v)
forall a b c. (a -> b -> c) -> b -> a -> c
flip g (Term f v) -> Maybe (g (Term f v)) -> g (Term f v)
forall a. a -> Maybe a -> a
fromMaybe (Term f v -> Maybe (g (Term f v))
h Term f v
t) (g (Term f v) -> g (Term f v)) -> g (Term f v) -> g (Term f v)
forall a b. (a -> b) -> a -> b
$ case Term f v -> ABT f v
forall (f :: * -> * -> *) v. Term f v -> ABT f v
out Term f v
t of
  Abs v
x Term f v
e -> v -> Term f v -> Term f v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
TAbs v
x (Term f v -> Term f v) -> g (Term f v) -> g (Term f v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Term f v -> Maybe (g (Term f v))) -> Term f v -> g (Term f v)
forall (g :: * -> *) (f :: * -> * -> *) v.
(Applicative g, Bifoldable f, Traversable (f v), Var v) =>
(Term f v -> Maybe (g (Term f v))) -> Term f v -> g (Term f v)
visit Term f v -> Maybe (g (Term f v))
h Term f v
e
  Tm f v (Term f v)
body -> f v (Term f v) -> Term f v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
TTm (f v (Term f v) -> Term f v) -> g (f v (Term f v)) -> g (Term f v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Term f v -> g (Term f v)) -> f v (Term f v) -> g (f v (Term f v))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> f v a -> f (f v b)
traverse ((Term f v -> Maybe (g (Term f v))) -> Term f v -> g (Term f v)
forall (g :: * -> *) (f :: * -> * -> *) v.
(Applicative g, Bifoldable f, Traversable (f v), Var v) =>
(Term f v -> Maybe (g (Term f v))) -> Term f v -> g (Term f v)
visit Term f v -> Maybe (g (Term f v))
h) f v (Term f v)
body

visitPure ::
  (Bifoldable f, Traversable (f v), Var v) =>
  (Term f v -> Maybe (Term f v)) ->
  Term f v ->
  Term f v
visitPure :: forall (f :: * -> * -> *) v.
(Bifoldable f, Traversable (f v), Var v) =>
(Term f v -> Maybe (Term f v)) -> Term f v -> Term f v
visitPure Term f v -> Maybe (Term f v)
h = Identity (Term f v) -> Term f v
forall a. Identity a -> a
runIdentity (Identity (Term f v) -> Term f v)
-> (Term f v -> Identity (Term f v)) -> Term f v -> Term f v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term f v -> Maybe (Identity (Term f v)))
-> Term f v -> Identity (Term f v)
forall (g :: * -> *) (f :: * -> * -> *) v.
(Applicative g, Bifoldable f, Traversable (f v), Var v) =>
(Term f v -> Maybe (g (Term f v))) -> Term f v -> g (Term f v)
visit ((Term f v -> Identity (Term f v))
-> Maybe (Term f v) -> Maybe (Identity (Term f v))
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Term f v -> Identity (Term f v)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Term f v) -> Maybe (Identity (Term f v)))
-> (Term f v -> Maybe (Term f v))
-> Term f v
-> Maybe (Identity (Term f v))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term f v -> Maybe (Term f v)
h)