{-# LANGUAGE DerivingVia #-}

module GardGround.Utils.ScopeMask where

-- this mostly a direct copy of https://github.com/AndrasKovacs/smalltt/blob/b5fc9b3747dcbf8c71ca9f976afb7b574b630996/src/LvlSet.hs
-- see LICENSE.smalltt for license information

import Data.Bits
import Data.Foldable (foldl')

newtype ScopeMask = ScopeMask Integer deriving (Eq, Bits) via Integer

instance Semigroup ScopeMask where
  (<>) = (.|.)
  {-# INLINE (<>) #-}

instance Monoid ScopeMask where
  mempty = ScopeMask 0
  {-# INLINE mempty #-}

insert :: Int -> ScopeMask -> ScopeMask
insert x (ScopeMask s) = ScopeMask (unsafeShiftL 1 x .|. s)
{-# INLINE insert #-}

member :: Int -> ScopeMask -> Bool
member x (ScopeMask s) = (unsafeShiftL 1 x .&. s) /= 0
{-# INLINE member #-}

toList :: ScopeMask -> [Int]
toList s' = reverse $ intern 0 s'
  where
    -- | `intern offset smask`
    intern :: Int -> ScopeMask -> [Int]
    intern _ (ScopeMask 0) = []
    intern i (ScopeMask s) =
      let nxt = intern (i + 1) (ScopeMask $ unsafeShiftR s 1) in
      if 1 .&. s /= 0 then i:nxt else nxt

fromList :: [Int] -> ScopeMask
fromList = foldl' (flip insert) mempty

instance Show ScopeMask where
  show = show . toList