{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeOperators #-}
module Data.VarMap (
  VarMap,
  empty,
  insert,
  delete,
  TypedIdx(..),
  lookup,
  disjointUnion,
  sink1,
  unsink1,
  subMap,
  superMap,
) 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))

disjointUnion :: Ord k => VarMap k env -> VarMap k env -> VarMap k env
disjointUnion (VarMap off1 cl1 m1) (VarMap off2 cl2 m2) | off1 == off2 =
  VarMap off1 (min cl1 cl2) (Map.unionWith (error "VarMap.disjointUnion: overlapping keys") m1 m2)
disjointUnion vm1 vm2 = disjointUnion (cleanup vm1) (cleanup vm2)

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)

superMap :: Eq k => Subenv env env' -> VarMap k env' -> VarMap k env
superMap subenv =
  let loop :: Subenv env env' -> Int -> [Int]
      loop SETop _ = []
      loop (SEYes sub) i = i : loop sub (i+1)
      loop (SENo sub) i = loop sub (i+1)

      newIndices = VS.fromList $ loop subenv 0
      modify off (k, (ty, i))
        | i + off < 0 = Nothing
        | i + off >= VS.length newIndices = error "VarMap.superMap: found negative indices in map"
        | otherwise = let j = newIndices VS.! (i + off)
                      in if j == -1 then Nothing else Just (k, (ty, j))

  in \(VarMap off _ mp) -> VarMap 0 0 (Map.fromAscList . mapMaybe (modify off) . Map.toAscList $ mp)

maybeCleanup :: VarMap k env -> VarMap k env
maybeCleanup vm@(VarMap _ interval mp)
  | let sz = Map.size mp
  , sz > 0, 2 * interval >= 3 * sz
  = cleanup vm
maybeCleanup vm = vm

cleanup :: VarMap k env -> VarMap k env
cleanup (VarMap off _ mp) = VarMap 0 0 (Map.mapMaybe (\(t, i) -> if i + off >= 0 then Just (t, i + off) else Nothing) mp)

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)))