{-# LANGUAGE CPP
           , GADTs
           , DataKinds
           , PolyKinds
           , FlexibleContexts
           , DeriveDataTypeable
           , ExistentialQuantification
           , TypeInType
           , UndecidableInstances
           , ScopedTypeVariables
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
----------------------------------------------------------------
--                                                    2016.04.28
-- |
-- Module      :  Language.Hakaru.Syntax.Variable
-- Copyright   :  Copyright (c) 2016 the Hakaru team
-- License     :  BSD3
-- Maintainer  :  wren@community.haskell.org
-- Stability   :  experimental
-- Portability :  GHC-only
--
-- An implementation of variables, for use with "Language.Hakaru.Syntax.ABT".
----------------------------------------------------------------
module Language.Hakaru.Syntax.Variable
    (
    -- * Our basic notion of variables.
      Variable(..)
    , varEq
    , VarEqTypeError(..)
    -- ** Variables with existentially quantified types
    , KindOf
    , SomeVariable(..)
    
    -- * Some helper types for \"heaps\", \"environments\", etc
    -- ** Typing environments; aka: sets of (typed) variables
    , VarSet(..)
    , emptyVarSet
    , singletonVarSet
    , fromVarSet
    , toVarSet
    , toVarSet1
    , varSetKeys
    , insertVarSet
    , deleteVarSet
    , memberVarSet
    , unionVarSet
    , intersectVarSet
    , sizeVarSet
    , nextVarID
    -- ** Substitutions; aka: maps from variables to their definitions
    , Assoc(..)
    , Assocs(..) -- TODO: hide the data constructors
    , emptyAssocs
    , singletonAssocs
    , fromAssocs
    , toAssocs
    , toAssocs1
    , insertAssoc
    , insertOrReplaceAssoc
    , insertAssocs
    , lookupAssoc
    , adjustAssoc
    , mapAssocs
    ) where

import           Data.Proxy        (KProxy(..))
import           Data.Typeable     (Typeable)
import           Data.Kind
import           Data.Text         (Text)
import           Data.IntMap       (IntMap)
import qualified Data.IntMap       as IM
import           Data.Function     (on)
import           Control.Exception (Exception, throw)
#if __GLASGOW_HASKELL__ < 710
import           Data.Monoid       (Monoid(..))
#endif

#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup
#endif

import Data.Number.Nat
import Language.Hakaru.Syntax.IClasses
-- TODO: factor the definition of the 'Sing' type family out from
-- the instances, so that we can make our ABT stuff totally independent
-- of the definition of Hakaru's types.
import Language.Hakaru.Types.Sing

----------------------------------------------------------------
----------------------------------------------------------------
-- TODO: should we make this type abstract, or type-class it?

-- TODO: alas we need to keep the Sing in order to make 'subst'
-- typesafe... Is there any way to work around that? Maybe only
-- define substitution for well-typed ABTs (i.e., what we produce
-- via typechecking a plain ABT)? If we can manage to get rid of
-- the Sing, then 'binder' and 'multibinder' would become much
-- simpler. Alas, it looks like we also need it for 'inferType' to
-- be well-typed... How can we avoid that?
--
-- TODO: what are the overhead costs of storing a Sing? Would
-- it be cheaper to store the SingI dictionary (and a Proxy,
-- as necessary)?


-- | A variable is a triple of a unique identifier ('varID'), a
-- hint for how to display things to humans ('varHint'), and a type
-- ('varType'). Notably, the hint is only used for display purposes,
-- and the type is only used for typing purposes; thus, the 'Eq'
-- and 'Ord' instances only look at the unique identifier, completely
-- ignoring the other two components. However, the 'varEq' function
-- does take the type into consideration (but still ignores the
-- hint).
--
-- N.B., the unique identifier is lazy so that we can tie-the-knot
-- in 'binder'.
data Variable (a :: k) = Variable
    { Variable a -> Text
varHint :: {-# UNPACK #-} !Text
    , Variable a -> Nat
varID   :: Nat -- N.B., lazy!
    , Variable a -> Sing a
varType :: !(Sing a)
    }

-- TODO: instance Read (Variable a)

-- HACK: this requires UndecidableInstances
instance Show1 (Sing :: k -> Type) => Show1 (Variable :: k -> Type) where
    showsPrec1 :: Int -> Variable i -> ShowS
showsPrec1 Int
p (Variable Text
hint Nat
i Sing i
typ) =
        Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
9)
            ( String -> ShowS
showString String
"Variable "
            ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Text -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec  Int
11 Text
hint
            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 -> Nat -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec  Int
11 Nat
i
            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 -> Sing i -> ShowS
forall k (a :: k -> *) (i :: k). Show1 a => Int -> a i -> ShowS
showsPrec1 Int
11 Sing i
typ
            )

instance Show (Sing a) => Show (Variable a) where
    showsPrec :: Int -> Variable a -> ShowS
showsPrec Int
p (Variable Text
hint Nat
i Sing a
typ) =
        Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
9)
            ( String -> ShowS
showString String
"Variable "
            ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Text -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec  Int
11 Text
hint
            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 -> Nat -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec  Int
11 Nat
i
            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 -> Sing a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec  Int
11 Sing a
typ
            )

-- BUG: these may not be consistent with the interpretation chosen by 'varEq'
instance Eq1 Variable where
    eq1 :: Variable i -> Variable i -> Bool
eq1 = Nat -> Nat -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Nat -> Nat -> Bool)
-> (Variable i -> Nat) -> Variable i -> Variable i -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` Variable i -> Nat
forall k (a :: k). Variable a -> Nat
varID

instance Eq (Variable a) where
    == :: Variable a -> Variable a -> Bool
(==) = Nat -> Nat -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Nat -> Nat -> Bool)
-> (Variable a -> Nat) -> Variable a -> Variable a -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID


-- BUG: this must be consistent with the 'Eq' instance, but should
-- also be consistent with the 'varEq' interpretation. In particular,
-- it's not clear how to make any Ord instance consistent with
-- interpretation #1 (unless we have some sort of `jmCompare` on
-- types!)
instance Ord (Variable a) where
    compare :: Variable a -> Variable a -> Ordering
compare = Nat -> Nat -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Nat -> Nat -> Ordering)
-> (Variable a -> Nat) -> Variable a -> Variable a -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID


-- TODO: so long as we don't go with interpretation #1 (because
-- that'd cause consistency issues with the 'Ord' instance) we could
-- simply use this to give a 'JmEq1' instance. Would help to minimize
-- the number of distinct concepts floating around...
--
-- | Compare to variables at possibly-different types. If the
-- variables are \"equal\", then they must in fact have the same
-- type. N.B., it is not entirely specified what this function
-- /means/ when two variables have the same 'varID' but different
-- 'varType'. However, so long as we use this function everywhere,
-- at least we'll be consistent.
--
-- Possible interpretations:
--
-- * We could /assume/ that when the 'varType's do not match the
-- variables are not equal. Upside: we can statically guarantee
-- that every variable is \"well-typed\" (by fiat). Downside: every
-- type has its own variable namespace, which is very confusing.
-- Also, the @Ord SomeVariable@ instance will be really difficult
-- to get right.
--
-- * We could /require/ that whenever two 'varID's match, their
-- 'varType's must also match. Upside: a single variable namespace.
-- Downside: if the types do not in fact match (e.g., the preprocessing
-- step for ensuring variable uniqueness is buggy), then we must
-- throw (or return) an 'VarEqTypeError' exception.
--
-- * We could /assert/ that whenever two 'varID's match, their
-- 'varType's must also match. Upsides: we get a single variable
-- namespace, and we get /O(1)/ equality checking. Downsides: if
-- the types do not in fact match, we'll probably segfault.
--
-- Whichever interpretation we choose, we must make sure that typing
-- contexts, binding environments, and so on all behave consistently.
varEq
    :: (Show1 (Sing :: k -> Type), JmEq1 (Sing :: k -> Type))
    => Variable (a :: k)
    -> Variable (b :: k)
    -> Maybe (TypeEq a b)
{-
-- Interpretation #1:
varEq x y =
    case jmEq1 (varType x) (varType y) of
    Just Refl | x == y -> Just Refl
    _                  -> Nothing
-}
-- Interpretation #2:
varEq :: Variable a -> Variable b -> Maybe (TypeEq a b)
varEq Variable a
x Variable b
y
    | Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x Nat -> Nat -> Bool
forall a. Eq a => a -> a -> Bool
== Variable b -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable b
y =
        case Sing a -> Sing b -> Maybe (TypeEq a b)
forall k (a :: k -> *) (i :: k) (j :: k).
JmEq1 a =>
a i -> a j -> Maybe (TypeEq i j)
jmEq1 (Variable a -> Sing a
forall k (a :: k). Variable a -> Sing a
varType Variable a
x) (Variable b -> Sing b
forall k (a :: k). Variable a -> Sing a
varType Variable b
y) of
        Just TypeEq a b
Refl -> TypeEq a a -> Maybe (TypeEq a a)
forall a. a -> Maybe a
Just TypeEq a a
forall k (a :: k). TypeEq a a
Refl
        Maybe (TypeEq a b)
Nothing   -> VarEqTypeError -> Maybe (TypeEq a b)
forall a e. Exception e => e -> a
throw (Variable a -> Variable b -> VarEqTypeError
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Variable b -> VarEqTypeError
VarEqTypeError Variable a
x Variable b
y)
    | Bool
otherwise = Maybe (TypeEq a b)
forall a. Maybe a
Nothing
{-
-- Interpretation #3:
varEq x y
    | varID x == varID y = Just (unsafeCoerce Refl)
    | otherwise          = Nothing
-}


-- TODO: is there any reason we ought to parameterize 'VarEqTypeError'
-- by the kind of the variables it closes over? Packaging up the
-- dictionaries seems fine for the 'Show' and 'Exception' instances,
-- but maybe elsewhere?
--
-- | An exception type for if we need to throw an error when two
-- variables do not have an equal 'varType'. This is mainly used
-- when 'varEq' chooses the second interpretation.
data VarEqTypeError where
    VarEqTypeError
        :: (Show1 (Sing :: k -> Type), JmEq1 (Sing :: k -> Type))
        => {-# UNPACK #-} !(Variable (a :: k))
        -> {-# UNPACK #-} !(Variable (b :: k))
        -> VarEqTypeError
    deriving (Typeable)

instance Show VarEqTypeError where
    showsPrec :: Int -> VarEqTypeError -> ShowS
showsPrec Int
p (VarEqTypeError Variable a
x Variable b
y) =
        Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
9)
            ( String -> ShowS
showString String
"VarEqTypeError "
            ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Variable a -> ShowS
forall k (a :: k -> *) (i :: k). Show1 a => Int -> a i -> ShowS
showsPrec1 Int
11 Variable a
x
            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 -> Variable b -> ShowS
forall k (a :: k -> *) (i :: k). Show1 a => Int -> a i -> ShowS
showsPrec1 Int
11 Variable b
y
            )

instance Exception VarEqTypeError


----------------------------------------------------------------
-- TODO: switch to using 'Some1' itself? Maybe no longer a good idea, due to the need for the kind parameter...

-- | Hide an existentially quantified parameter to 'Variable'.
--
-- Because the 'Variable' type is poly-kinded, we need to be careful
-- not to erase too much type\/kind information. Thus, we parameterize
-- the 'SomeVariable' type by the /kind/ of the type we existentially
-- quantify over. This is necessary for giving 'Eq' and 'Ord'
-- instances since we can only compare variables whose types live
-- in the same kind.
--
-- N.B., the 'Ord' instance assumes that 'varEq' uses either the
-- second or third interpretation. If 'varEq' uses the first
-- interpretation then, the 'Eq' instance (which uses 'varEq') will
-- be inconsistent with the 'Ord' instance!
data SomeVariable (kproxy :: KProxy k) =
    forall (a :: k) . SomeVariable
        {-# UNPACK #-} !(Variable (a :: k))


-- | Convenient synonym to refer to the kind of a type variable:
-- @type KindOf (a :: k) = ('KProxy :: KProxy k)@
type KindOf (a :: k) = ('KProxy :: KProxy k)


-- This instance requires the 'JmEq1' and 'Show1' constraints because we use 'varEq'.
instance (JmEq1 (Sing :: k -> Type), Show1 (Sing :: k -> Type))
    => Eq (SomeVariable (kproxy :: KProxy k))
    where
    SomeVariable Variable a
x == :: SomeVariable kproxy -> SomeVariable kproxy -> Bool
== SomeVariable Variable a
y =
        case Variable a -> Variable a -> Maybe (TypeEq a a)
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Variable b -> Maybe (TypeEq a b)
varEq Variable a
x Variable a
y of
        Just TypeEq a a
Refl -> Bool
True
        Maybe (TypeEq a a)
Nothing   -> Bool
False


-- This instance requires the 'JmEq1' and 'Show1' constraints because 'Ord' requires the 'Eq' instance, which in turn requires those constraints.
instance (JmEq1 (Sing :: k -> Type), Show1 (Sing :: k -> Type))
    => Ord (SomeVariable (kproxy :: KProxy k))
    where
    SomeVariable Variable a
x compare :: SomeVariable kproxy -> SomeVariable kproxy -> Ordering
`compare` SomeVariable Variable a
y =
        Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x Nat -> Nat -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
y


-- TODO: instance Read SomeVariable


instance Show1 (Sing :: k -> Type)
    => Show (SomeVariable (kproxy :: KProxy k))
    where
    showsPrec :: Int -> SomeVariable kproxy -> ShowS
showsPrec Int
p (SomeVariable Variable a
v) =
        Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
9)
            ( String -> ShowS
showString String
"SomeVariable "
            ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Variable a -> ShowS
forall k (a :: k -> *) (i :: k). Show1 a => Int -> a i -> ShowS
showsPrec1 Int
11 Variable a
v
            )


----------------------------------------------------------------
-- | A set of (typed) variables.
newtype VarSet (kproxy :: KProxy k) =
    VarSet { VarSet kproxy -> IntMap (SomeVariable kproxy)
unVarSet :: IntMap (SomeVariable kproxy) }

instance Show1 (Sing :: k -> Type) => Show (VarSet (kproxy :: KProxy k)) where
    showsPrec :: Int -> VarSet kproxy -> ShowS
showsPrec Int
p (VarSet IntMap (SomeVariable kproxy)
xs) =
        Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
9)
            ( String -> ShowS
showString String
"VarSet "
            ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> IntMap (SomeVariable kproxy) -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec  Int
11 IntMap (SomeVariable kproxy)
xs
            )

instance (Eq (SomeVariable (kproxy :: KProxy k))) => Eq (VarSet kproxy) where
  VarSet IntMap (SomeVariable kproxy)
s1 == :: VarSet kproxy -> VarSet kproxy -> Bool
== VarSet IntMap (SomeVariable kproxy)
s2 = IntMap (SomeVariable kproxy)
s1 IntMap (SomeVariable kproxy)
-> IntMap (SomeVariable kproxy) -> Bool
forall a. Eq a => a -> a -> Bool
== IntMap (SomeVariable kproxy)
s2

-- | Return the successor of the largest 'varID' of all the variables
-- in the set. Thus, we return zero for the empty set and non-zero
-- for non-empty sets.
nextVarID :: VarSet kproxy -> Nat
nextVarID :: VarSet kproxy -> Nat
nextVarID (VarSet IntMap (SomeVariable kproxy)
xs)
    | IntMap (SomeVariable kproxy) -> Bool
forall a. IntMap a -> Bool
IM.null IntMap (SomeVariable kproxy)
xs = Nat
0
    | Bool
otherwise  =
        case IntMap (SomeVariable kproxy) -> (Int, SomeVariable kproxy)
forall a. IntMap a -> (Int, a)
IM.findMax IntMap (SomeVariable kproxy)
xs of
        (Int
_, SomeVariable Variable a
x) -> Nat
1 Nat -> Nat -> Nat
forall a. Num a => a -> a -> a
+ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x


emptyVarSet :: VarSet kproxy
emptyVarSet :: VarSet kproxy
emptyVarSet = IntMap (SomeVariable kproxy) -> VarSet kproxy
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy) -> VarSet kproxy
VarSet IntMap (SomeVariable kproxy)
forall a. IntMap a
IM.empty

singletonVarSet :: Variable a -> VarSet (KindOf a)
singletonVarSet :: Variable a -> VarSet (KindOf a)
singletonVarSet Variable a
x =
    IntMap (SomeVariable (KindOf a)) -> VarSet (KindOf a)
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy) -> VarSet kproxy
VarSet (IntMap (SomeVariable (KindOf a)) -> VarSet (KindOf a))
-> IntMap (SomeVariable (KindOf a)) -> VarSet (KindOf a)
forall a b. (a -> b) -> a -> b
$ Int -> SomeVariable (KindOf a) -> IntMap (SomeVariable (KindOf a))
forall a. Int -> a -> IntMap a
IM.singleton (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x) (Variable a -> SomeVariable (KindOf a)
forall k (kproxy :: KProxy k) (a :: k).
Variable a -> SomeVariable kproxy
SomeVariable Variable a
x)

fromVarSet :: VarSet kproxy -> [SomeVariable kproxy]
fromVarSet :: VarSet kproxy -> [SomeVariable kproxy]
fromVarSet (VarSet IntMap (SomeVariable kproxy)
xs) = IntMap (SomeVariable kproxy) -> [SomeVariable kproxy]
forall a. IntMap a -> [a]
IM.elems IntMap (SomeVariable kproxy)
xs

-- | Convert a list of variables into a variable set.
--
-- In the event that multiple variables have conflicting 'varID',
-- the latter variable will be kept. This generally won't matter
-- because we're treating the list as a /set/. In the cases where
-- it would matter, chances are you're going to encounter a
-- 'VarEqTypeError' sooner or later anyways.
toVarSet :: [SomeVariable kproxy] -> VarSet kproxy
toVarSet :: [SomeVariable kproxy] -> VarSet kproxy
toVarSet = IntMap (SomeVariable kproxy) -> VarSet kproxy
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy) -> VarSet kproxy
VarSet (IntMap (SomeVariable kproxy) -> VarSet kproxy)
-> ([SomeVariable kproxy] -> IntMap (SomeVariable kproxy))
-> [SomeVariable kproxy]
-> VarSet kproxy
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (SomeVariable kproxy)
-> [SomeVariable kproxy] -> IntMap (SomeVariable kproxy)
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy)
-> [SomeVariable kproxy] -> IntMap (SomeVariable kproxy)
go IntMap (SomeVariable kproxy)
forall a. IntMap a
IM.empty
    where
    go :: IntMap (SomeVariable kproxy)
-> [SomeVariable kproxy] -> IntMap (SomeVariable kproxy)
go IntMap (SomeVariable kproxy)
vars [SomeVariable kproxy]
_ | IntMap (SomeVariable kproxy)
vars IntMap (SomeVariable kproxy) -> Bool -> Bool
`seq` Bool
False = String -> IntMap (SomeVariable kproxy)
forall a. HasCallStack => String -> a
error String
"toVarSet: the impossible happened"
    go IntMap (SomeVariable kproxy)
vars []     = IntMap (SomeVariable kproxy)
vars
    go IntMap (SomeVariable kproxy)
vars (SomeVariable kproxy
x:[SomeVariable kproxy]
xs) = IntMap (SomeVariable kproxy)
-> [SomeVariable kproxy] -> IntMap (SomeVariable kproxy)
go (Int
-> SomeVariable kproxy
-> IntMap (SomeVariable kproxy)
-> IntMap (SomeVariable kproxy)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ SomeVariable kproxy -> Nat
forall k (kproxy :: KProxy k). SomeVariable kproxy -> Nat
someVarID SomeVariable kproxy
x) SomeVariable kproxy
x IntMap (SomeVariable kproxy)
vars) [SomeVariable kproxy]
xs

    someVarID :: SomeVariable kproxy -> Nat
    someVarID :: SomeVariable kproxy -> Nat
someVarID (SomeVariable Variable a
x) = Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x


-- | Convert a list of variables into a variable set.
--
-- In the event that multiple variables have conflicting 'varID',
-- the latter variable will be kept. This generally won't matter
-- because we're treating the list as a /set/. In the cases where
-- it would matter, chances are you're going to encounter a
-- 'VarEqTypeError' sooner or later anyways.
toVarSet1 :: List1 Variable (xs :: [k]) -> VarSet (kproxy :: KProxy k)
toVarSet1 :: List1 Variable xs -> VarSet kproxy
toVarSet1 = [SomeVariable kproxy] -> VarSet kproxy
forall k (kproxy :: KProxy k).
[SomeVariable kproxy] -> VarSet kproxy
toVarSet ([SomeVariable kproxy] -> VarSet kproxy)
-> (List1 Variable xs -> [SomeVariable kproxy])
-> List1 Variable xs
-> VarSet kproxy
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List1 Variable xs -> [SomeVariable kproxy]
forall k (xs :: [k]) (kproxy :: KProxy k).
List1 Variable xs -> [SomeVariable kproxy]
someVariables
    where
    -- N.B., this conversion maintains the variable ordering.
    someVariables
        :: List1 Variable (xs :: [k])
        -> [SomeVariable (kproxy :: KProxy k)]
    someVariables :: List1 Variable xs -> [SomeVariable kproxy]
someVariables List1 Variable xs
Nil1         = []
    someVariables (Cons1 Variable x
x List1 Variable xs
xs) = Variable x -> SomeVariable kproxy
forall k (kproxy :: KProxy k) (a :: k).
Variable a -> SomeVariable kproxy
SomeVariable Variable x
x SomeVariable kproxy
-> [SomeVariable kproxy] -> [SomeVariable kproxy]
forall a. a -> [a] -> [a]
: List1 Variable xs -> [SomeVariable kproxy]
forall k (xs :: [k]) (kproxy :: KProxy k).
List1 Variable xs -> [SomeVariable kproxy]
someVariables List1 Variable xs
xs

instance Semigroup (VarSet kproxy) where
    VarSet IntMap (SomeVariable kproxy)
xs <> :: VarSet kproxy -> VarSet kproxy -> VarSet kproxy
<> VarSet IntMap (SomeVariable kproxy)
ys = IntMap (SomeVariable kproxy) -> VarSet kproxy
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy) -> VarSet kproxy
VarSet (IntMap (SomeVariable kproxy)
-> IntMap (SomeVariable kproxy) -> IntMap (SomeVariable kproxy)
forall a. IntMap a -> IntMap a -> IntMap a
IM.union IntMap (SomeVariable kproxy)
xs IntMap (SomeVariable kproxy)
ys) -- TODO: remove bias; crash if conflicting definitions

instance Monoid (VarSet kproxy) where
    mempty :: VarSet kproxy
mempty  = VarSet kproxy
forall k (kproxy :: KProxy k). VarSet kproxy
emptyVarSet
#if !(MIN_VERSION_base(4,11,0))
    mappend = (<>)
#endif
    mconcat :: [VarSet kproxy] -> VarSet kproxy
mconcat = IntMap (SomeVariable kproxy) -> VarSet kproxy
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy) -> VarSet kproxy
VarSet (IntMap (SomeVariable kproxy) -> VarSet kproxy)
-> ([VarSet kproxy] -> IntMap (SomeVariable kproxy))
-> [VarSet kproxy]
-> VarSet kproxy
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [IntMap (SomeVariable kproxy)] -> IntMap (SomeVariable kproxy)
forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions ([IntMap (SomeVariable kproxy)] -> IntMap (SomeVariable kproxy))
-> ([VarSet kproxy] -> [IntMap (SomeVariable kproxy)])
-> [VarSet kproxy]
-> IntMap (SomeVariable kproxy)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarSet kproxy -> IntMap (SomeVariable kproxy))
-> [VarSet kproxy] -> [IntMap (SomeVariable kproxy)]
forall a b. (a -> b) -> [a] -> [b]
map VarSet kproxy -> IntMap (SomeVariable kproxy)
forall k (kproxy :: KProxy k).
VarSet kproxy -> IntMap (SomeVariable kproxy)
unVarSet

varSetKeys :: VarSet a -> [Int]
varSetKeys :: VarSet a -> [Int]
varSetKeys (VarSet IntMap (SomeVariable a)
set) = IntMap (SomeVariable a) -> [Int]
forall a. IntMap a -> [Int]
IM.keys IntMap (SomeVariable a)
set

insertVarSet :: Variable a -> VarSet (KindOf a) -> VarSet (KindOf a)
insertVarSet :: Variable a -> VarSet (KindOf a) -> VarSet (KindOf a)
insertVarSet Variable a
x (VarSet IntMap (SomeVariable (KindOf a))
xs) =
    case
        (Int
 -> SomeVariable (KindOf a)
 -> SomeVariable (KindOf a)
 -> SomeVariable (KindOf a))
-> Int
-> SomeVariable (KindOf a)
-> IntMap (SomeVariable (KindOf a))
-> (Maybe (SomeVariable (KindOf a)),
    IntMap (SomeVariable (KindOf a)))
forall a.
(Int -> a -> a -> a) -> Int -> a -> IntMap a -> (Maybe a, IntMap a)
IM.insertLookupWithKey
            (\Int
_ SomeVariable (KindOf a)
v' SomeVariable (KindOf a)
_ -> SomeVariable (KindOf a)
v')
            (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x)
            (Variable a -> SomeVariable (KindOf a)
forall k (kproxy :: KProxy k) (a :: k).
Variable a -> SomeVariable kproxy
SomeVariable Variable a
x)
            IntMap (SomeVariable (KindOf a))
xs
    of
    (Maybe (SomeVariable (KindOf a))
Nothing, IntMap (SomeVariable (KindOf a))
xs') -> IntMap (SomeVariable (KindOf a)) -> VarSet (KindOf a)
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy) -> VarSet kproxy
VarSet IntMap (SomeVariable (KindOf a))
xs'
    (Just SomeVariable (KindOf a)
_,  IntMap (SomeVariable (KindOf a))
_)   -> String -> VarSet (KindOf a)
forall a. HasCallStack => String -> a
error String
"insertVarSet: variable is already assigned!"


deleteVarSet :: Variable a -> VarSet (KindOf a) -> VarSet (KindOf a)
deleteVarSet :: Variable a -> VarSet (KindOf a) -> VarSet (KindOf a)
deleteVarSet Variable a
x (VarSet IntMap (SomeVariable (KindOf a))
xs) =
    --- BUG: use some sort of deleteLookupWithKey to make sure we got the right one...
    IntMap (SomeVariable (KindOf a)) -> VarSet (KindOf a)
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy) -> VarSet kproxy
VarSet (IntMap (SomeVariable (KindOf a)) -> VarSet (KindOf a))
-> IntMap (SomeVariable (KindOf a)) -> VarSet (KindOf a)
forall a b. (a -> b) -> a -> b
$ Int
-> IntMap (SomeVariable (KindOf a))
-> IntMap (SomeVariable (KindOf a))
forall a. Int -> IntMap a -> IntMap a
IM.delete (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x) IntMap (SomeVariable (KindOf a))
xs


memberVarSet
    :: (Show1 (Sing :: k -> Type), JmEq1 (Sing :: k -> Type))
    => Variable (a :: k)
    -> VarSet (kproxy :: KProxy k)
    -> Bool
memberVarSet :: Variable a -> VarSet kproxy -> Bool
memberVarSet Variable a
x (VarSet IntMap (SomeVariable kproxy)
xs) =
    -- HACK: can't use do-notation here for GADT reasons
    case Int -> IntMap (SomeVariable kproxy) -> Maybe (SomeVariable kproxy)
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x) IntMap (SomeVariable kproxy)
xs of
    Maybe (SomeVariable kproxy)
Nothing                -> Bool
False
    Just (SomeVariable Variable a
x') -> 
        case Variable a -> Variable a -> Maybe (TypeEq a a)
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Variable b -> Maybe (TypeEq a b)
varEq Variable a
x Variable a
x' of
        Maybe (TypeEq a a)
Nothing -> Bool
False
        Just TypeEq a a
_  -> Bool
True

-- NB: The union and intersection operations are left biased.
-- What is the best behaviour when we have two variables with
-- different types in the set?
unionVarSet
    :: forall k (kproxy :: KProxy k)
    .  (Show1 (Sing :: k -> Type), JmEq1 (Sing :: k -> Type))
    => VarSet kproxy
    -> VarSet kproxy
    -> VarSet kproxy
unionVarSet :: VarSet kproxy -> VarSet kproxy -> VarSet kproxy
unionVarSet (VarSet IntMap (SomeVariable kproxy)
s1) (VarSet IntMap (SomeVariable kproxy)
s2) = IntMap (SomeVariable kproxy) -> VarSet kproxy
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy) -> VarSet kproxy
VarSet (IntMap (SomeVariable kproxy)
-> IntMap (SomeVariable kproxy) -> IntMap (SomeVariable kproxy)
forall a. IntMap a -> IntMap a -> IntMap a
IM.union IntMap (SomeVariable kproxy)
s1 IntMap (SomeVariable kproxy)
s2)

intersectVarSet
    :: forall k (kproxy :: KProxy k)
    .  (Show1 (Sing :: k -> Type), JmEq1 (Sing :: k -> Type))
    => VarSet kproxy
    -> VarSet kproxy
    -> VarSet kproxy
intersectVarSet :: VarSet kproxy -> VarSet kproxy -> VarSet kproxy
intersectVarSet (VarSet IntMap (SomeVariable kproxy)
s1) (VarSet IntMap (SomeVariable kproxy)
s2) = IntMap (SomeVariable kproxy) -> VarSet kproxy
forall k (kproxy :: KProxy k).
IntMap (SomeVariable kproxy) -> VarSet kproxy
VarSet (IntMap (SomeVariable kproxy)
-> IntMap (SomeVariable kproxy) -> IntMap (SomeVariable kproxy)
forall a b. IntMap a -> IntMap b -> IntMap a
IM.intersection IntMap (SomeVariable kproxy)
s1 IntMap (SomeVariable kproxy)
s2)

sizeVarSet :: VarSet a -> Int
sizeVarSet :: VarSet a -> Int
sizeVarSet (VarSet IntMap (SomeVariable a)
xs) = IntMap (SomeVariable a) -> Int
forall a. IntMap a -> Int
IM.size IntMap (SomeVariable a)
xs

----------------------------------------------------------------
-- BUG: haddock doesn't like annotations on GADT constructors. So
-- here we'll avoid using the GADT syntax, even though it'd make
-- the data type declaration prettier\/cleaner.
-- <https://github.com/hakaru-dev/hakaru/issues/6>
--
-- | A pair of variable and term, both of the same Hakaru type.
data Assoc (ast :: k -> Type)
    = forall (a :: k) . Assoc
        {-# UNPACK #-} !(Variable a)
        !(ast a)

instance (Show1 (Sing :: k -> Type), Show1 (ast :: k -> Type))
    => Show (Assoc ast)
    where
    showsPrec :: Int -> Assoc ast -> ShowS
showsPrec Int
p (Assoc Variable a
x ast a
e) =
        Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
9)
            ( String -> ShowS
showString String
"Assoc "
            ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Variable a -> ShowS
forall k (a :: k -> *) (i :: k). Show1 a => Int -> a i -> ShowS
showsPrec1 Int
11 Variable a
x
            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 -> ast a -> ShowS
forall k (a :: k -> *) (i :: k). Show1 a => Int -> a i -> ShowS
showsPrec1 Int
11 ast a
e
            )


-- BUG: since multiple 'varEq'-distinct variables could have the
-- same ID, we should really have the elements be a list of
-- associations (or something more efficient; e.g., if 'Sing' is
-- hashable).
--
-- | A set of variable\/term associations.
--
-- N.B., the current implementation assumes 'varEq' uses either the
-- second or third interpretations; that is, it is impossible to
-- have a single 'varID' be shared by multiple variables (i.e., at
-- different types). If you really want the first interpretation,
-- then the implementation must be updated.
newtype Assocs ast = Assocs { Assocs ast -> IntMap (Assoc ast)
unAssocs :: IntMap (Assoc ast) }

instance (Show1 (Sing :: k -> Type), Show1 (ast :: k -> Type))
    => Show (Assocs ast)
    where
    showsPrec :: Int -> Assocs ast -> ShowS
showsPrec Int
p Assocs ast
rho =
        Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
9)
            ( String -> ShowS
showString String
"toAssocs "
            ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Assoc ast -> ShowS) -> [Assoc ast] -> ShowS
forall a. (a -> ShowS) -> [a] -> ShowS
showListWith Assoc ast -> ShowS
forall a. Show a => a -> ShowS
shows (Assocs ast -> [Assoc ast]
forall k (ast :: k -> *). Assocs ast -> [Assoc ast]
fromAssocs Assocs ast
rho)
            )

-- | The empty set of associations.
emptyAssocs :: Assocs abt
emptyAssocs :: Assocs abt
emptyAssocs = IntMap (Assoc abt) -> Assocs abt
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs IntMap (Assoc abt)
forall a. IntMap a
IM.empty

-- | A single association.
singletonAssocs :: Variable a -> f a -> Assocs f
singletonAssocs :: Variable a -> f a -> Assocs f
singletonAssocs Variable a
x f a
e =
    IntMap (Assoc f) -> Assocs f
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs (IntMap (Assoc f) -> Assocs f) -> IntMap (Assoc f) -> Assocs f
forall a b. (a -> b) -> a -> b
$ Int -> Assoc f -> IntMap (Assoc f)
forall a. Int -> a -> IntMap a
IM.singleton (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x) (Variable a -> f a -> Assoc f
forall k (ast :: k -> *) (a :: k). Variable a -> ast a -> Assoc ast
Assoc Variable a
x f a
e)

-- | Convert an association list into a list of associations.
fromAssocs :: Assocs ast -> [Assoc ast]
fromAssocs :: Assocs ast -> [Assoc ast]
fromAssocs (Assocs IntMap (Assoc ast)
rho) = IntMap (Assoc ast) -> [Assoc ast]
forall a. IntMap a -> [a]
IM.elems IntMap (Assoc ast)
rho

-- | Convert a list of associations into an association list. In
-- the event of conflict, later associations override earlier ones.
toAssocs :: [Assoc ast] -> Assocs ast
toAssocs :: [Assoc ast] -> Assocs ast
toAssocs = IntMap (Assoc ast) -> Assocs ast
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs (IntMap (Assoc ast) -> Assocs ast)
-> ([Assoc ast] -> IntMap (Assoc ast)) -> [Assoc ast] -> Assocs ast
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IntMap (Assoc ast) -> Assoc ast -> IntMap (Assoc ast))
-> IntMap (Assoc ast) -> [Assoc ast] -> IntMap (Assoc ast)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl IntMap (Assoc ast) -> Assoc ast -> IntMap (Assoc ast)
forall k (ast :: k -> *).
IntMap (Assoc ast) -> Assoc ast -> IntMap (Assoc ast)
step IntMap (Assoc ast)
forall a. IntMap a
IM.empty
    where
    step :: IntMap (Assoc ast) -> Assoc ast -> IntMap (Assoc ast)
    step :: IntMap (Assoc ast) -> Assoc ast -> IntMap (Assoc ast)
step IntMap (Assoc ast)
xes xe :: Assoc ast
xe@(Assoc Variable a
x ast a
_) = Int -> Assoc ast -> IntMap (Assoc ast) -> IntMap (Assoc ast)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x) Assoc ast
xe IntMap (Assoc ast)
xes


-- TODO: Do we also want a zipped curried variant: @List1 (Pair1 Variable ast) xs@?
-- | Convert an unzipped list of curried associations into an
-- association list. In the event of conflict, later associations
-- override earlier ones.
toAssocs1 :: List1 Variable xs -> List1 ast xs -> Assocs ast
toAssocs1 :: List1 Variable xs -> List1 ast xs -> Assocs ast
toAssocs1 = \List1 Variable xs
xs List1 ast xs
es -> IntMap (Assoc ast) -> Assocs ast
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs (IntMap (Assoc ast)
-> List1 Variable xs -> List1 ast xs -> IntMap (Assoc ast)
forall k (ast :: k -> *) (xs :: [k]).
IntMap (Assoc ast)
-> List1 Variable xs -> List1 ast xs -> IntMap (Assoc ast)
go IntMap (Assoc ast)
forall a. IntMap a
IM.empty List1 Variable xs
xs List1 ast xs
es)
    where
    go  :: IntMap (Assoc ast)
        -> List1 Variable xs
        -> List1 ast xs
        -> IntMap (Assoc ast)
    -- BUG: GHC claims the patterns are non-exhaustive here
    go :: IntMap (Assoc ast)
-> List1 Variable xs -> List1 ast xs -> IntMap (Assoc ast)
go IntMap (Assoc ast)
m List1 Variable xs
Nil1         List1 ast xs
Nil1         = IntMap (Assoc ast)
m
    go IntMap (Assoc ast)
m (Cons1 Variable x
x List1 Variable xs
xs) (Cons1 ast x
e List1 ast xs
es) =
        IntMap (Assoc ast)
-> List1 Variable xs -> List1 ast xs -> IntMap (Assoc ast)
forall k (ast :: k -> *) (xs :: [k]).
IntMap (Assoc ast)
-> List1 Variable xs -> List1 ast xs -> IntMap (Assoc ast)
go (Int -> Assoc ast -> IntMap (Assoc ast) -> IntMap (Assoc ast)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable x -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable x
x) (Variable x -> ast x -> Assoc ast
forall k (ast :: k -> *) (a :: k). Variable a -> ast a -> Assoc ast
Assoc Variable x
x ast x
ast x
e) IntMap (Assoc ast)
m) List1 Variable xs
xs List1 ast xs
List1 ast xs
es

instance Semigroup (Assocs abt) where
    Assocs IntMap (Assoc abt)
xs <> :: Assocs abt -> Assocs abt -> Assocs abt
<> Assocs IntMap (Assoc abt)
ys = IntMap (Assoc abt) -> Assocs abt
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs (IntMap (Assoc abt) -> IntMap (Assoc abt) -> IntMap (Assoc abt)
forall a. IntMap a -> IntMap a -> IntMap a
IM.union IntMap (Assoc abt)
xs IntMap (Assoc abt)
ys) -- TODO: remove bias; crash if conflicting definitions

instance Monoid (Assocs abt) where
    mempty :: Assocs abt
mempty  = Assocs abt
forall k (abt :: k -> *). Assocs abt
emptyAssocs
#if !(MIN_VERSION_base(4,11,0))
    mappend = (<>)
#endif
    mconcat :: [Assocs abt] -> Assocs abt
mconcat = IntMap (Assoc abt) -> Assocs abt
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs (IntMap (Assoc abt) -> Assocs abt)
-> ([Assocs abt] -> IntMap (Assoc abt))
-> [Assocs abt]
-> Assocs abt
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [IntMap (Assoc abt)] -> IntMap (Assoc abt)
forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions ([IntMap (Assoc abt)] -> IntMap (Assoc abt))
-> ([Assocs abt] -> [IntMap (Assoc abt)])
-> [Assocs abt]
-> IntMap (Assoc abt)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Assocs abt -> IntMap (Assoc abt))
-> [Assocs abt] -> [IntMap (Assoc abt)]
forall a b. (a -> b) -> [a] -> [b]
map Assocs abt -> IntMap (Assoc abt)
forall k (ast :: k -> *). Assocs ast -> IntMap (Assoc ast)
unAssocs


-- If we actually do have a list (etc) of variables for each ID,
-- and want to add the new binding to whatever old ones, then it
-- looks like there's no way to do that in one pass of both the
-- IntMap and the list.
--
-- | Add an association to the set of associations.
--
-- HACK: if the variable is already associated with some term then
-- we throw an error! In the future it'd be better to take some
-- sort of continuation to decide between (a) replacing the old
-- binding, (b) throwing an exception, or (c) safely wrapping the
-- result up with 'Maybe'
insertAssoc :: Assoc ast -> Assocs ast -> Assocs ast
insertAssoc :: Assoc ast -> Assocs ast -> Assocs ast
insertAssoc v :: Assoc ast
v@(Assoc Variable a
x ast a
_) (Assocs IntMap (Assoc ast)
xs) =
    case (Int -> Assoc ast -> Assoc ast -> Assoc ast)
-> Int
-> Assoc ast
-> IntMap (Assoc ast)
-> (Maybe (Assoc ast), IntMap (Assoc ast))
forall a.
(Int -> a -> a -> a) -> Int -> a -> IntMap a -> (Maybe a, IntMap a)
IM.insertLookupWithKey (\Int
_ Assoc ast
v' Assoc ast
_ -> Assoc ast
v') (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x) Assoc ast
v IntMap (Assoc ast)
xs of
    (Maybe (Assoc ast)
Nothing, IntMap (Assoc ast)
xs') -> IntMap (Assoc ast) -> Assocs ast
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs IntMap (Assoc ast)
xs'
    (Just Assoc ast
_,  IntMap (Assoc ast)
_  ) -> String -> Assocs ast
forall a. HasCallStack => String -> a
error String
"insertAssoc: variable is already assigned!"

insertOrReplaceAssoc :: Assoc ast -> Assocs ast -> Assocs ast
insertOrReplaceAssoc :: Assoc ast -> Assocs ast -> Assocs ast
insertOrReplaceAssoc v :: Assoc ast
v@(Assoc Variable a
x ast a
_) (Assocs IntMap (Assoc ast)
xs) =
    IntMap (Assoc ast) -> Assocs ast
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs (IntMap (Assoc ast) -> Assocs ast)
-> IntMap (Assoc ast) -> Assocs ast
forall a b. (a -> b) -> a -> b
$ Int -> Assoc ast -> IntMap (Assoc ast) -> IntMap (Assoc ast)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x) Assoc ast
v IntMap (Assoc ast)
xs

insertAssocs :: Assocs ast -> Assocs ast -> Assocs ast
insertAssocs :: Assocs ast -> Assocs ast -> Assocs ast
insertAssocs (Assocs IntMap (Assoc ast)
from) Assocs ast
to = (Assoc ast -> Assocs ast -> Assocs ast)
-> Assocs ast -> IntMap (Assoc ast) -> Assocs ast
forall a b. (a -> b -> b) -> b -> IntMap a -> b
IM.foldr Assoc ast -> Assocs ast -> Assocs ast
forall k (ast :: k -> *). Assoc ast -> Assocs ast -> Assocs ast
insertAssoc Assocs ast
to IntMap (Assoc ast)
from

-- | Adjust an association so existing variable refers to different
-- value. Does nothing if variable not present.
adjustAssoc :: Variable (a :: k)
            -> (Assoc ast -> Assoc ast)
            -> Assocs ast
            -> Assocs ast
adjustAssoc :: Variable a -> (Assoc ast -> Assoc ast) -> Assocs ast -> Assocs ast
adjustAssoc Variable a
x Assoc ast -> Assoc ast
f (Assocs IntMap (Assoc ast)
xs) =
    IntMap (Assoc ast) -> Assocs ast
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs (IntMap (Assoc ast) -> Assocs ast)
-> IntMap (Assoc ast) -> Assocs ast
forall a b. (a -> b) -> a -> b
$ (Assoc ast -> Assoc ast)
-> Int -> IntMap (Assoc ast) -> IntMap (Assoc ast)
forall a. (a -> a) -> Int -> IntMap a -> IntMap a
IM.adjust Assoc ast -> Assoc ast
f (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x) IntMap (Assoc ast)
xs

-- | Look up a variable and return the associated term.
--
-- N.B., this function is robust to all interpretations of 'varEq'.
lookupAssoc
    :: (Show1 (Sing :: k -> Type), JmEq1 (Sing :: k -> Type))
    => Variable (a :: k)
    -> Assocs ast
    -> Maybe (ast a)
lookupAssoc :: Variable a -> Assocs ast -> Maybe (ast a)
lookupAssoc Variable a
x (Assocs IntMap (Assoc ast)
xs) = do
    Assoc Variable a
x' ast a
e' <- Int -> IntMap (Assoc ast) -> Maybe (Assoc ast)
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Nat -> Int
fromNat (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x) IntMap (Assoc ast)
xs
    TypeEq a a
Refl        <- Variable a -> Variable a -> Maybe (TypeEq a a)
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Variable b -> Maybe (TypeEq a b)
varEq Variable a
x Variable a
x'
    ast a -> Maybe (ast a)
forall (m :: * -> *) a. Monad m => a -> m a
return ast a
e'
{-
-- for @Assocs abt = IntMap [Assoc abt]@ this should work:
lookupAssoc x (Assocs xss) =
    go x <$> IM.lookup (fromNat $ varID x) xss
    where
    go x []                 = Nothing
    go x (Assoc x' e' : xs) =
        case varEq x x' of
        Just Refl -> Just e'
        Nothing   -> go x xs
-}

mapAssocs :: (Assoc ast1 -> Assoc ast2) -> Assocs ast1 -> Assocs ast2
mapAssocs :: (Assoc ast1 -> Assoc ast2) -> Assocs ast1 -> Assocs ast2
mapAssocs Assoc ast1 -> Assoc ast2
f (Assocs IntMap (Assoc ast1)
xs) = IntMap (Assoc ast2) -> Assocs ast2
forall k (ast :: k -> *). IntMap (Assoc ast) -> Assocs ast
Assocs ((Assoc ast1 -> Assoc ast2)
-> IntMap (Assoc ast1) -> IntMap (Assoc ast2)
forall a b. (a -> b) -> IntMap a -> IntMap b
IM.map Assoc ast1 -> Assoc ast2
f IntMap (Assoc ast1)
xs)
                

----------------------------------------------------------------
----------------------------------------------------------- fin.