-- | Maps from de Bruijn indices, implemented using 'IntMap's from de
-- Bruijn levels.

module Agda.Utils.IndexMap
  (IndexMap, invariant, nil, snoc, append, fromList, update, index)
  where

import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import qualified Data.List as L

import Agda.Utils.Impossible

-- | Maps from de Bruijn indices.

-- Note that the map goes from de Bruijn *levels*.
--
-- Invariant: The map keys form a (possibly empty) contiguous range
-- from 0 upwards.

newtype IndexMap a = IndexMap { forall a. IndexMap a -> IntMap a
indexMap :: IntMap a }
  deriving IndexMap a -> IndexMap a -> Bool
(IndexMap a -> IndexMap a -> Bool)
-> (IndexMap a -> IndexMap a -> Bool) -> Eq (IndexMap a)
forall a. Eq a => IndexMap a -> IndexMap a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => IndexMap a -> IndexMap a -> Bool
== :: IndexMap a -> IndexMap a -> Bool
$c/= :: forall a. Eq a => IndexMap a -> IndexMap a -> Bool
/= :: IndexMap a -> IndexMap a -> Bool
Eq

instance Show a => Show (IndexMap a) where
  show :: IndexMap a -> String
show !(IndexMap IntMap a
m) =
    String
"Agda.Utils.IndexMap.fromList " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [a] -> String
forall a. Show a => a -> String
show (IntMap a -> [a]
forall a. IntMap a -> [a]
IntMap.elems IntMap a
m)

-- | The 'IndexMap' invariant.

invariant :: IndexMap a -> Bool
invariant :: forall a. IndexMap a -> Bool
invariant !(IndexMap IntMap a
m) = IntMap a -> [Int]
forall a. IntMap a -> [Int]
IntMap.keys IntMap a
m [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== [Int
0 .. IntMap a -> Int
forall a. IntMap a -> Int
IntMap.size IntMap a
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

-- | Converts a de Bruijn index to a map index.

mapIndex :: IndexMap a -> Int -> Int
mapIndex :: forall a. IndexMap a -> Int -> Int
mapIndex !IndexMap a
m !Int
i = IntMap a -> Int
forall a. IntMap a -> Int
IntMap.size (IndexMap a -> IntMap a
forall a. IndexMap a -> IntMap a
indexMap IndexMap a
m) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

-- | Is the index in range for the given map?

inRange :: Int -> IndexMap a -> Bool
inRange :: forall a. Int -> IndexMap a -> Bool
inRange !Int
i !IndexMap a
m = Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< IntMap a -> Int
forall a. IntMap a -> Int
IntMap.size (IndexMap a -> IntMap a
forall a. IndexMap a -> IntMap a
indexMap IndexMap a
m)

-- | An empty map.

nil :: IndexMap a
nil :: forall a. IndexMap a
nil = IntMap a -> IndexMap a
forall a. IntMap a -> IndexMap a
IndexMap IntMap a
forall a. IntMap a
IntMap.empty

-- | Inserts a new binding, which will have de Bruijn index 0, and
-- shifts other bindings one step.

snoc :: IndexMap a -> a -> IndexMap a
snoc :: forall a. IndexMap a -> a -> IndexMap a
snoc !(IndexMap IntMap a
m) !a
x = IntMap a -> IndexMap a
forall a. IntMap a -> IndexMap a
IndexMap (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert (IntMap a -> Int
forall a. IntMap a -> Int
IntMap.size IntMap a
m) a
x IntMap a
m)

-- | Inserts a number of new bindings. The last element in the list
-- will have de Bruijn index 0.

append :: IndexMap a -> [a] -> IndexMap a
append :: forall a. IndexMap a -> [a] -> IndexMap a
append !IndexMap a
m ![a]
xs = (IndexMap a -> a -> IndexMap a) -> IndexMap a -> [a] -> IndexMap a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' IndexMap a -> a -> IndexMap a
forall a. IndexMap a -> a -> IndexMap a
snoc IndexMap a
m [a]
xs

-- | Converts a list to an 'IndexMap'. The last element in the list
-- will have de Bruijn index 0.

fromList :: [a] -> IndexMap a
fromList :: forall a. [a] -> IndexMap a
fromList = IndexMap a -> [a] -> IndexMap a
forall a. IndexMap a -> [a] -> IndexMap a
append IndexMap a
forall a. IndexMap a
nil

-- Overwrites a binding with the given de Bruijn index.
--
-- Precondition: The binding must already exist in the map.

update :: Int -> a -> IndexMap a -> IndexMap a
update :: forall a. Int -> a -> IndexMap a -> IndexMap a
update !Int
i !a
x !IndexMap a
m
  | Int -> IndexMap a -> Bool
forall a. Int -> IndexMap a -> Bool
inRange Int
i IndexMap a
m = IntMap a -> IndexMap a
forall a. IntMap a -> IndexMap a
IndexMap (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert (IndexMap a -> Int -> Int
forall a. IndexMap a -> Int -> Int
mapIndex IndexMap a
m Int
i) a
x (IndexMap a -> IntMap a
forall a. IndexMap a -> IntMap a
indexMap IndexMap a
m))
  | Bool
otherwise   = IndexMap a
forall a. HasCallStack => a
__IMPOSSIBLE__

-- Looks up the thing with the given de Bruijn index.
--
-- Precondition: The index should be in range.

index :: Int -> IndexMap a -> a
index :: forall a. Int -> IndexMap a -> a
index !Int
i !IndexMap a
m = case Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup (IndexMap a -> Int -> Int
forall a. IndexMap a -> Int -> Int
mapIndex IndexMap a
m Int
i) (IndexMap a -> IntMap a
forall a. IndexMap a -> IntMap a
indexMap IndexMap a
m) of
  Just a
x  -> a
x
  Maybe a
Nothing -> a
forall a. HasCallStack => a
__IMPOSSIBLE__