{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeOperators #-}
module Data.VarMap (
  VarMap,
  empty,
  insert,
  delete,
  TypedIdx(..),
  lookup,
  sink1,
  unsink1,
  subMap,
) where

import Prelude hiding (lookup)

import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import Data.Maybe (mapMaybe)
import Data.Some
import qualified Data.Vector.Storable as VS
import Unsafe.Coerce

import AST.Env
import AST.Types
import AST.Weaken


type role VarMap _ nominal  -- ensure that 'env' is not phantom
data VarMap k (env :: [Ty]) =
  VarMap Int  -- ^ Global offset; must be added to any value in the map in order to get the proper index
         Int  -- ^ Time since last cleanup
         (Map k (Some STy, Int))
deriving instance Show k => Show (VarMap k env)

empty :: VarMap k env
empty = VarMap 0 0 Map.empty

insert :: Ord k => k -> STy t -> Idx env t -> VarMap k env -> VarMap k env
insert k ty idx (VarMap off interval mp) =
  maybeCleanup $ VarMap off (interval + 1) (Map.insert k (Some ty, idx2int idx - off) mp)

delete :: Ord k => k -> VarMap k env -> VarMap k env
delete k (VarMap off interval mp) =
  maybeCleanup $ VarMap off (interval + 1) (Map.delete k mp)

data TypedIdx env t = TypedIdx (STy t) (Idx env t)
  deriving (Show)

lookup :: Ord k => k -> VarMap k env -> Maybe (Some (TypedIdx env))
lookup k (VarMap off _ mp) = do
  (Some ty, i) <- Map.lookup k mp
  idx <- unsafeInt2idx (i + off)
  return (Some (TypedIdx ty idx))

sink1 :: VarMap k env -> VarMap k (t : env)
sink1 (VarMap off interval mp) = VarMap (off + 1) interval mp

unsink1 :: VarMap k (t : env) -> VarMap k env
unsink1 (VarMap off interval mp) = VarMap (off - 1) interval mp

subMap :: Eq k => Subenv env env' -> VarMap k env -> VarMap k env'
subMap subenv =
  let bools = let loop :: Subenv env env' -> [Bool]
                  loop SETop = []
                  loop (SEYes sub) = True : loop sub
                  loop (SENo sub) = False : loop sub
              in VS.fromList $ loop subenv
      newIndices = VS.init $ VS.scanl' (\n b -> if b then n + 1 else n) (0 :: Int) bools
      modify off (k, (ty, i))
        | i + off < 0 = Nothing
        | i + off >= VS.length bools = error "VarMap.subMap: found negative indices in map"
        | bools VS.! (i + off) = Just (k, (ty, newIndices VS.! (i + off)))
        | otherwise = Nothing
  in \(VarMap off _ mp) -> VarMap 0 0 (Map.fromAscList . mapMaybe (modify off) . Map.toAscList $ mp)

maybeCleanup :: VarMap k env -> VarMap k env
maybeCleanup (VarMap off interval mp)
  | let sz = Map.size mp
  , sz > 0, 2 * interval >= 3 * sz
  = VarMap off 0 (Map.filter (\(_, i) -> i + off >= 0) mp)
maybeCleanup vm = vm

unsafeInt2idx :: Int -> Maybe (Idx env t)
unsafeInt2idx = \n -> if n < 0 then Nothing else Just (go n)
  where
    go :: Int -> Idx env t
    go 0 = unsafeCoerce IZ
    go n = unsafeCoerce (IS (go (n-1)))