content: 
  module AOC.Common.Intcode.Memory (
      MonadMem(..)
    , Memory(..)
    , mRegLens
    , MemRef(..)
    , initMemRef
    , freezeMemRef
    ) where
  
  import           Control.DeepSeq
  import           Control.Lens
  import           Control.Monad.Except
  import           Control.Monad.Primitive
  import           Control.Monad.Reader
  import           Control.Monad.State
  import           Data.Conduino
  import           Data.Generics.Labels         ()
  import           Data.Map                     (Map)
  import           Data.Primitive.MutVar
  import           GHC.Generics
  import           Numeric.Natural              (Natural)
  import qualified Data.Map                     as M
  import qualified Data.Vector.Storable         as VS
  import qualified Data.Vector.Storable.Mutable as MVS
  
  
  class Monad m => MonadMem m where
      mRead      :: m Int
      mCurr      :: m Natural
      mPeek      :: Natural -> m Int
      mSeek      :: Natural -> m ()
      mWrite     :: Natural -> Int -> m ()
      mShiftBase :: Int -> m ()
      mWithBase  :: Int -> m Int
  
      -- mFreeze    :: m Memory
      -- mPutMem      :: Memory -> m ()
  
  
  data Memory = Mem
      { mPos  :: Natural
      , mBase :: Int
      , mRegs :: Map Natural Int
      }
    deriving (Eq, Ord, Show, Generic)
  instance NFData Memory
  
  instance Monad m => MonadMem (StateT Memory m) where
      mRead = do
        m@Mem{..} <- get
        M.findWithDefault 0 mPos mRegs <$ put (m { mPos = mPos + 1 })
      mCurr = gets mPos
      mPeek i = gets $ M.findWithDefault 0 i . mRegs
      mSeek z = modify $ \m -> m { mPos = z }
      mWrite i x = modify $ \m -> m { mRegs = M.insert i x (mRegs m) }
      mShiftBase b = modify $ \m -> m { mBase = mBase m + b }
      mWithBase i = gets $ (+ i) . mBase
  
      -- mFreeze = get
      -- mPutMem   = put
  
  instance MonadMem m => MonadMem (Pipe i o u m) where
      mRead = lift mRead
      mCurr = lift mCurr
      mPeek = lift . mPeek
      mSeek = lift . mSeek
      mWrite i = lift . mWrite i
      mShiftBase = lift . mShiftBase
      mWithBase  = lift . mWithBase
      -- mFreeze = lift mFreeze
      -- mPutMem   = lift . mPutMem
  
  instance MonadMem m => MonadMem (ExceptT e m) where
      mRead = lift mRead
      mCurr = lift mCurr
      mPeek = lift . mPeek
      mSeek = lift . mSeek
      mWrite i = lift . mWrite i
      mShiftBase = lift . mShiftBase
      mWithBase  = lift . mWithBase
      -- mFreeze = lift mFreeze
      -- mPutMem   = lift . mPutMem
  
  mRegLens :: Natural -> Lens' Memory Int
  mRegLens i = #mRegs . at i . non 0
  
  data MemRef s = MemRef
      { mrPos  :: MutVar s Natural
      , mrBase :: MutVar s Int
      , mrRegs :: MutVar s (VS.MVector s Int)
      }
  
  initMemRef :: (PrimMonad m, s ~ PrimState m) => Memory -> m (MemRef s)
  initMemRef Mem{..} = do
      mrPos  <- newMutVar mPos
      mrBase <- newMutVar mBase
      mrRegs <- case M.lookupMax mRegs of
        Nothing     -> newMutVar =<< MVS.new 0
        Just (n, _) -> do
          let r = VS.generate (fromIntegral n * 10 + 1) $ \i -> M.findWithDefault 0 (fromIntegral i) mRegs
          newMutVar =<< VS.thaw r
      pure MemRef{..}
  
  freezeMemRef :: (PrimMonad m, s ~ PrimState m) => MemRef s -> m Memory
  freezeMemRef MemRef{..} = do
      mPos  <- readMutVar mrPos
      mBase <- readMutVar mrBase
      mRegs <- fmap toRegs . VS.freeze =<< readMutVar mrRegs
      pure Mem{..}
    where
      toRegs = M.filter (/= 0) . M.fromList . zip [0..] . VS.toList
  
  
  instance (PrimMonad m, s ~ PrimState m) => MonadMem (ReaderT (MemRef s) m) where
      mRead = ask >>= \MemRef{..} -> do
        i <- fromIntegral <$> atomicModifyMutVar' mrPos (\i -> (i+1, i))
        mPeek i
      mCurr = readMutVar =<< asks mrPos
      mPeek i = do
        r <- readMutVar =<< asks mrRegs
        if i' < MVS.length r
          then MVS.unsafeRead r i'
          else pure 0
        where
          i' = fromIntegral i
      mSeek i = (`writeMutVar` i) =<< asks mrPos
      mWrite i x = ask >>= \MemRef{..} -> do
          r <- readMutVar mrRegs
          let l0 = MVS.length r
          if i' < MVS.length r
            then MVS.unsafeWrite r i' x
            -- else trace "grow" $ do
            else do
              let l1 = (i' + 1) * 2
              regs' <- MVS.unsafeGrow r (l1 - l0)
              forM_ [l0 .. l1 - 1] $ \j ->
                MVS.unsafeWrite regs' j 0
              MVS.unsafeWrite regs' i' x
              writeMutVar mrRegs regs'
        where
          i' = fromIntegral i
      mShiftBase b = (`modifyMutVar'` (+ b)) =<< asks mrBase
      mWithBase  i = fmap (+ i) . readMutVar =<< asks mrBase
  
      -- mFreeze = freezeMemRef =<< ask
      -- mPutMem Mem{..} = ask >>= \MemRef{..} -> do
      --   writeMutVar mrPos  mPos
      --   writeMutVar mrBase mBase
      --   case M.lookupMax mRegs of
      --     Nothing     -> writeMutVar mrRegs =<< MVS.new 0
      --     Just (n, _) -> do
      --       let r = VS.generate (fromIntegral n + 1) $ \i -> M.findWithDefault 0 (fromIntegral i) mRegs
      --       writeMutVar mrRegs =<< VS.thaw r
  
  