-- |
-- Module      :  Cryptol.Backend.SeqMap
-- Copyright   :  (c) 2013-2021 Galois, Inc.
-- License     :  BSD3
-- Maintainer  :  cryptol@galois.com
-- Stability   :  provisional
-- Portability :  portable

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DoAndIfThenElse #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}

module Cryptol.Backend.SeqMap
  ( -- * Sequence Maps
    SeqMap
  , indexSeqMap
  , lookupSeqMap
  , finiteSeqMap
  , infiniteSeqMap
  , enumerateSeqMap
  , streamSeqMap
  , reverseSeqMap
  , updateSeqMap
  , dropSeqMap
  , concatSeqMap
  , splitSeqMap
  , memoMap
  , delaySeqMap
  , zipSeqMap
  , mapSeqMap
  , mergeSeqMap
  , barrelShifter
  , shiftSeqByInteger

  , IndexSegment(..)
  ) where

import qualified Control.Exception as X
import Control.Monad
import Control.Monad.IO.Class
import Data.Bits
import Data.List
import Data.IORef
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import Cryptol.Backend
import Cryptol.Backend.Concrete (Concrete)
import Cryptol.Backend.Monad (Unsupported(..))

import Cryptol.TypeCheck.Solver.InfNat(Nat'(..))
import Cryptol.Utils.Panic

-- | A sequence map represents a mapping from nonnegative integer indices
--   to values.  These are used to represent both finite and infinite sequences.
data SeqMap sym a
  = IndexSeqMap  !(Integer -> SEval sym a)
  | UpdateSeqMap !(Map Integer (SEval sym a))
                 !(SeqMap sym a)
  | MemoSeqMap
     !Nat'
     !(IORef (Map Integer a))
     !(IORef (Integer -> SEval sym a))

indexSeqMap :: (Integer -> SEval sym a) -> SeqMap sym a
indexSeqMap :: (Integer -> SEval sym a) -> SeqMap sym a
indexSeqMap = (Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap

lookupSeqMap :: Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap :: SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap (IndexSeqMap Integer -> SEval sym a
f) Integer
i = Integer -> SEval sym a
f Integer
i
lookupSeqMap (UpdateSeqMap Map Integer (SEval sym a)
m SeqMap sym a
xs) Integer
i =
  case Integer -> Map Integer (SEval sym a) -> Maybe (SEval sym a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Integer
i Map Integer (SEval sym a)
m of
    Just SEval sym a
x  -> SEval sym a
x
    Maybe (SEval sym a)
Nothing -> SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
xs Integer
i
lookupSeqMap (MemoSeqMap Nat'
sz IORef (Map Integer a)
cache IORef (Integer -> SEval sym a)
eval) Integer
i =
  do Maybe a
mz <- IO (Maybe a) -> SEval sym (Maybe a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Integer -> Map Integer a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Integer
i (Map Integer a -> Maybe a) -> IO (Map Integer a) -> IO (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Map Integer a) -> IO (Map Integer a)
forall a. IORef a -> IO a
readIORef IORef (Map Integer a)
cache)
     case Maybe a
mz of
       Just a
z  -> a -> SEval sym a
forall (m :: * -> *) a. Monad m => a -> m a
return a
z
       Maybe a
Nothing ->
         do Integer -> SEval sym a
f <- IO (Integer -> SEval sym a) -> SEval sym (Integer -> SEval sym a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef (Integer -> SEval sym a) -> IO (Integer -> SEval sym a)
forall a. IORef a -> IO a
readIORef IORef (Integer -> SEval sym a)
eval)
            a
v <- Integer -> SEval sym a
f Integer
i
            Int
msz <- IO Int -> SEval sym Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> SEval sym Int) -> IO Int -> SEval sym Int
forall a b. (a -> b) -> a -> b
$ IORef (Map Integer a)
-> (Map Integer a -> (Map Integer a, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Map Integer a)
cache (\Map Integer a
m ->
                        let m' :: Map Integer a
m' = Integer -> a -> Map Integer a -> Map Integer a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Integer
i a
v Map Integer a
m in (Map Integer a
m', Map Integer a -> Int
forall k a. Map k a -> Int
Map.size Map Integer a
m'))
            -- If we memoize the entire map, overwrite the evaluation closure to let
            -- the garbage collector reap it
            Bool -> SEval sym () -> SEval sym ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (case Nat'
sz of Nat'
Inf -> Bool
False; Nat Integer
sz' -> Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
msz Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
sz')
                 (IO () -> SEval sym ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef (Integer -> SEval sym a) -> (Integer -> SEval sym a) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Integer -> SEval sym a)
eval
                    (\Integer
j -> String -> [String] -> SEval sym a
forall a. HasCallStack => String -> [String] -> a
panic String
"lookupSeqMap" [String
"Messed up size accounting", Nat' -> String
forall a. Show a => a -> String
show Nat'
sz, Integer -> String
forall a. Show a => a -> String
show Integer
j])))
            a -> SEval sym a
forall (m :: * -> *) a. Monad m => a -> m a
return a
v

instance Backend sym => Functor (SeqMap sym) where
  fmap :: (a -> b) -> SeqMap sym a -> SeqMap sym b
fmap a -> b
f SeqMap sym a
xs = (Integer -> SEval sym b) -> SeqMap sym b
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap (\Integer
i -> a -> b
f (a -> b) -> SEval sym a -> SEval sym b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
xs Integer
i)

-- | Generate a finite sequence map from a list of values
finiteSeqMap :: Backend sym => sym -> [SEval sym a] -> SeqMap sym a
finiteSeqMap :: sym -> [SEval sym a] -> SeqMap sym a
finiteSeqMap sym
sym [SEval sym a]
xs =
   Map Integer (SEval sym a) -> SeqMap sym a -> SeqMap sym a
forall sym a.
Map Integer (SEval sym a) -> SeqMap sym a -> SeqMap sym a
UpdateSeqMap
      ([(Integer, SEval sym a)] -> Map Integer (SEval sym a)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([Integer] -> [SEval sym a] -> [(Integer, SEval sym a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Integer
0..] [SEval sym a]
xs))
      ((Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap (\Integer
i -> sym -> Integer -> SEval sym a
forall sym a. Backend sym => sym -> Integer -> SEval sym a
invalidIndex sym
sym Integer
i))

-- | Generate an infinite sequence map from a stream of values
infiniteSeqMap :: Backend sym => sym -> [SEval sym a] -> SEval sym (SeqMap sym a)
infiniteSeqMap :: sym -> [SEval sym a] -> SEval sym (SeqMap sym a)
infiniteSeqMap sym
sym [SEval sym a]
xs =
   -- TODO: use an int-trie?
   sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
forall sym a.
Backend sym =>
sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
memoMap sym
sym Nat'
Inf ((Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i -> [SEval sym a] -> Integer -> SEval sym a
forall i a. Integral i => [a] -> i -> a
genericIndex [SEval sym a]
xs Integer
i)

-- | Create a finite list of length @n@ of the values from @[0..n-1]@ in
--   the given the sequence emap.
enumerateSeqMap :: (Backend sym, Integral n) => n -> SeqMap sym a -> [SEval sym a]
enumerateSeqMap :: n -> SeqMap sym a -> [SEval sym a]
enumerateSeqMap n
n SeqMap sym a
m = [ SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
m  Integer
i | Integer
i <- [Integer
0 .. (n -> Integer
forall a. Integral a => a -> Integer
toInteger n
n)Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1] ]

-- | Create an infinite stream of all the values in a sequence map
streamSeqMap :: Backend sym => SeqMap sym a -> [SEval sym a]
streamSeqMap :: SeqMap sym a -> [SEval sym a]
streamSeqMap SeqMap sym a
m = [ SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
m Integer
i | Integer
i <- [Integer
0..] ]

-- | Reverse the order of a finite sequence map
reverseSeqMap :: Backend sym =>
  Integer {- ^ Size of the sequence map -} ->
  SeqMap sym a ->
  SeqMap sym a
reverseSeqMap :: Integer -> SeqMap sym a -> SeqMap sym a
reverseSeqMap Integer
n SeqMap sym a
vals = (Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i -> SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
vals (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
i)

updateSeqMap :: SeqMap sym a -> Integer -> SEval sym a -> SeqMap sym a
updateSeqMap :: SeqMap sym a -> Integer -> SEval sym a -> SeqMap sym a
updateSeqMap (UpdateSeqMap Map Integer (SEval sym a)
m SeqMap sym a
sm) Integer
i SEval sym a
x = Map Integer (SEval sym a) -> SeqMap sym a -> SeqMap sym a
forall sym a.
Map Integer (SEval sym a) -> SeqMap sym a -> SeqMap sym a
UpdateSeqMap (Integer
-> SEval sym a
-> Map Integer (SEval sym a)
-> Map Integer (SEval sym a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Integer
i SEval sym a
x Map Integer (SEval sym a)
m) SeqMap sym a
sm
updateSeqMap SeqMap sym a
xs Integer
i SEval sym a
x = Map Integer (SEval sym a) -> SeqMap sym a -> SeqMap sym a
forall sym a.
Map Integer (SEval sym a) -> SeqMap sym a -> SeqMap sym a
UpdateSeqMap (Integer -> SEval sym a -> Map Integer (SEval sym a)
forall k a. k -> a -> Map k a
Map.singleton Integer
i SEval sym a
x) SeqMap sym a
xs

-- | Concatenate the first @n@ values of the first sequence map onto the
--   beginning of the second sequence map.
concatSeqMap :: Backend sym => Integer -> SeqMap sym a -> SeqMap sym a -> SeqMap sym a
concatSeqMap :: Integer -> SeqMap sym a -> SeqMap sym a -> SeqMap sym a
concatSeqMap Integer
n SeqMap sym a
x SeqMap sym a
y =
    (Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i ->
       if Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
n
         then SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
x Integer
i
         else SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
y (Integer
iInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
n)

-- | Given a number @n@ and a sequence map, return two new sequence maps:
--   the first containing the values from @[0..n-1]@ and the next containing
--   the values from @n@ onward.
splitSeqMap :: Backend sym => Integer -> SeqMap sym a -> (SeqMap sym a, SeqMap sym a)
splitSeqMap :: Integer -> SeqMap sym a -> (SeqMap sym a, SeqMap sym a)
splitSeqMap Integer
n SeqMap sym a
xs = (SeqMap sym a
hd,SeqMap sym a
tl)
  where
  hd :: SeqMap sym a
hd = SeqMap sym a
xs
  tl :: SeqMap sym a
tl = (Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i -> SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
xs (Integer
iInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
n)

-- | Drop the first @n@ elements of the given 'SeqMap'.
dropSeqMap :: Backend sym => Integer -> SeqMap sym a -> SeqMap sym a
dropSeqMap :: Integer -> SeqMap sym a -> SeqMap sym a
dropSeqMap Integer
0 SeqMap sym a
xs = SeqMap sym a
xs
dropSeqMap Integer
n SeqMap sym a
xs = (Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i -> SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
xs (Integer
iInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
n)

delaySeqMap :: Backend sym => sym -> SEval sym (SeqMap sym a) -> SEval sym (SeqMap sym a)
delaySeqMap :: sym -> SEval sym (SeqMap sym a) -> SEval sym (SeqMap sym a)
delaySeqMap sym
sym SEval sym (SeqMap sym a)
xs =
  do SEval sym (SeqMap sym a)
xs' <- sym
-> SEval sym (SeqMap sym a) -> SEval sym (SEval sym (SeqMap sym a))
forall sym a.
Backend sym =>
sym -> SEval sym a -> SEval sym (SEval sym a)
sDelay sym
sym SEval sym (SeqMap sym a)
xs
     SeqMap sym a -> SEval sym (SeqMap sym a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SeqMap sym a -> SEval sym (SeqMap sym a))
-> SeqMap sym a -> SEval sym (SeqMap sym a)
forall a b. (a -> b) -> a -> b
$ (Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i -> do SeqMap sym a
m <- SEval sym (SeqMap sym a)
xs'; SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
m Integer
i

-- | Given a sequence map, return a new sequence map that is memoized using
--   a finite map memo table.
memoMap :: Backend sym => sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)

-- Sequence is alreay memoized, just return it
memoMap :: sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
memoMap sym
_sym Nat'
_sz x :: SeqMap sym a
x@(MemoSeqMap{}) = SeqMap sym a -> SEval sym (SeqMap sym a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure SeqMap sym a
x

memoMap sym
sym Nat'
sz SeqMap sym a
x = do
  CallStack
stk <- sym -> SEval sym CallStack
forall sym. Backend sym => sym -> SEval sym CallStack
sGetCallStack sym
sym
  IORef (Map Integer a)
cache <- IO (IORef (Map Integer a)) -> SEval sym (IORef (Map Integer a))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef (Map Integer a)) -> SEval sym (IORef (Map Integer a)))
-> IO (IORef (Map Integer a)) -> SEval sym (IORef (Map Integer a))
forall a b. (a -> b) -> a -> b
$ Map Integer a -> IO (IORef (Map Integer a))
forall a. a -> IO (IORef a)
newIORef (Map Integer a -> IO (IORef (Map Integer a)))
-> Map Integer a -> IO (IORef (Map Integer a))
forall a b. (a -> b) -> a -> b
$ Map Integer a
forall k a. Map k a
Map.empty
  IORef (Integer -> SEval sym a)
evalRef <- IO (IORef (Integer -> SEval sym a))
-> SEval sym (IORef (Integer -> SEval sym a))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef (Integer -> SEval sym a))
 -> SEval sym (IORef (Integer -> SEval sym a)))
-> IO (IORef (Integer -> SEval sym a))
-> SEval sym (IORef (Integer -> SEval sym a))
forall a b. (a -> b) -> a -> b
$ (Integer -> SEval sym a) -> IO (IORef (Integer -> SEval sym a))
forall a. a -> IO (IORef a)
newIORef ((Integer -> SEval sym a) -> IO (IORef (Integer -> SEval sym a)))
-> (Integer -> SEval sym a) -> IO (IORef (Integer -> SEval sym a))
forall a b. (a -> b) -> a -> b
$ CallStack -> Integer -> SEval sym a
eval CallStack
stk
  SeqMap sym a -> SEval sym (SeqMap sym a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Nat'
-> IORef (Map Integer a)
-> IORef (Integer -> SEval sym a)
-> SeqMap sym a
forall sym a.
Nat'
-> IORef (Map Integer a)
-> IORef (Integer -> SEval sym a)
-> SeqMap sym a
MemoSeqMap Nat'
sz IORef (Map Integer a)
cache IORef (Integer -> SEval sym a)
evalRef)

  where
    eval :: CallStack -> Integer -> SEval sym a
eval CallStack
stk Integer
i = sym -> CallStack -> SEval sym a -> SEval sym a
forall sym a.
Backend sym =>
sym -> CallStack -> SEval sym a -> SEval sym a
sWithCallStack sym
sym CallStack
stk (SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
x Integer
i)


-- | Apply the given evaluation function pointwise to the two given
--   sequence maps.
zipSeqMap ::
  Backend sym =>
  sym ->
  (a -> a -> SEval sym a) ->
  Nat' ->
  SeqMap sym a ->
  SeqMap sym a ->
  SEval sym (SeqMap sym a)
zipSeqMap :: sym
-> (a -> a -> SEval sym a)
-> Nat'
-> SeqMap sym a
-> SeqMap sym a
-> SEval sym (SeqMap sym a)
zipSeqMap sym
sym a -> a -> SEval sym a
f Nat'
sz SeqMap sym a
x SeqMap sym a
y =
  sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
forall sym a.
Backend sym =>
sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
memoMap sym
sym Nat'
sz ((Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i -> SEval sym (SEval sym a) -> SEval sym a
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (a -> a -> SEval sym a
f (a -> a -> SEval sym a)
-> SEval sym a -> SEval sym (a -> SEval sym a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
x Integer
i SEval sym (a -> SEval sym a)
-> SEval sym a -> SEval sym (SEval sym a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
y Integer
i))

-- | Apply the given function to each value in the given sequence map
mapSeqMap ::
  Backend sym =>
  sym ->
  (a -> SEval sym a) ->
  Nat' ->
  SeqMap sym a ->
  SEval sym (SeqMap sym a)
mapSeqMap :: sym
-> (a -> SEval sym a)
-> Nat'
-> SeqMap sym a
-> SEval sym (SeqMap sym a)
mapSeqMap sym
sym a -> SEval sym a
f Nat'
sz SeqMap sym a
x =
  sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
forall sym a.
Backend sym =>
sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
memoMap sym
sym Nat'
sz ((Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i -> a -> SEval sym a
f (a -> SEval sym a) -> SEval sym a -> SEval sym a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
x Integer
i)


{-# INLINE mergeSeqMap #-}
mergeSeqMap :: Backend sym =>
  sym ->
  (SBit sym -> a -> a -> SEval sym a) ->
  SBit sym ->
  SeqMap sym a ->
  SeqMap sym a ->
  SeqMap sym a
mergeSeqMap :: sym
-> (SBit sym -> a -> a -> SEval sym a)
-> SBit sym
-> SeqMap sym a
-> SeqMap sym a
-> SeqMap sym a
mergeSeqMap sym
sym SBit sym -> a -> a -> SEval sym a
f SBit sym
c SeqMap sym a
x SeqMap sym a
y =
  (Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
IndexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i -> sym
-> (SBit sym -> a -> a -> SEval sym a)
-> SBit sym
-> SEval sym a
-> SEval sym a
-> SEval sym a
forall sym a.
Backend sym =>
sym
-> (SBit sym -> a -> a -> SEval sym a)
-> SBit sym
-> SEval sym a
-> SEval sym a
-> SEval sym a
mergeEval sym
sym SBit sym -> a -> a -> SEval sym a
f SBit sym
c (SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
x Integer
i) (SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
y Integer
i)


{-# INLINE shiftSeqByInteger #-}
shiftSeqByInteger :: Backend sym =>
  sym ->
  (SBit sym -> a -> a -> SEval sym a)
     {- ^ if/then/else operation of values -} ->
  (Integer -> Integer -> Maybe Integer)
     {- ^ reindexing operation -} ->
  SEval sym a {- ^ zero value -} ->
  Nat' {- ^ size of the sequence -} ->
  SeqMap sym a {- ^ sequence to shift -} ->
  SInteger sym {- ^ shift amount, assumed to be in range [0,len] -} ->
  SEval sym (SeqMap sym a)
shiftSeqByInteger :: sym
-> (SBit sym -> a -> a -> SEval sym a)
-> (Integer -> Integer -> Maybe Integer)
-> SEval sym a
-> Nat'
-> SeqMap sym a
-> SInteger sym
-> SEval sym (SeqMap sym a)
shiftSeqByInteger sym
sym SBit sym -> a -> a -> SEval sym a
merge Integer -> Integer -> Maybe Integer
reindex SEval sym a
zro Nat'
m SeqMap sym a
xs SInteger sym
idx
  | Just Integer
j <- sym -> SInteger sym -> Maybe Integer
forall sym. Backend sym => sym -> SInteger sym -> Maybe Integer
integerAsLit sym
sym SInteger sym
idx = SeqMap sym a -> Integer -> SEval sym (SeqMap sym a)
shiftOp SeqMap sym a
xs Integer
j
  | Bool
otherwise =
      do (Integer
n, [SBit sym]
idx_bits) <- sym -> Nat' -> SInteger sym -> SEval sym (Integer, [SBit sym])
forall sym.
Backend sym =>
sym -> Nat' -> SInteger sym -> SEval sym (Integer, [SBit sym])
enumerateIntBits sym
sym Nat'
m SInteger sym
idx
         sym
-> (SBit sym -> a -> a -> SEval sym a)
-> (SeqMap sym a -> Integer -> SEval sym (SeqMap sym a))
-> Nat'
-> SeqMap sym a
-> Integer
-> [IndexSegment sym]
-> SEval sym (SeqMap sym a)
forall sym a.
Backend sym =>
sym
-> (SBit sym -> a -> a -> SEval sym a)
-> (SeqMap sym a -> Integer -> SEval sym (SeqMap sym a))
-> Nat'
-> SeqMap sym a
-> Integer
-> [IndexSegment sym]
-> SEval sym (SeqMap sym a)
barrelShifter sym
sym SBit sym -> a -> a -> SEval sym a
merge SeqMap sym a -> Integer -> SEval sym (SeqMap sym a)
shiftOp Nat'
m SeqMap sym a
xs Integer
n ((SBit sym -> IndexSegment sym) -> [SBit sym] -> [IndexSegment sym]
forall a b. (a -> b) -> [a] -> [b]
map SBit sym -> IndexSegment sym
forall sym. SBit sym -> IndexSegment sym
BitIndexSegment [SBit sym]
idx_bits)
 where
   shiftOp :: SeqMap sym a -> Integer -> SEval sym (SeqMap sym a)
shiftOp SeqMap sym a
vs Integer
shft =
     SeqMap sym a -> SEval sym (SeqMap sym a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SeqMap sym a -> SEval sym (SeqMap sym a))
-> SeqMap sym a -> SEval sym (SeqMap sym a)
forall a b. (a -> b) -> a -> b
$ (Integer -> SEval sym a) -> SeqMap sym a
forall sym a. (Integer -> SEval sym a) -> SeqMap sym a
indexSeqMap ((Integer -> SEval sym a) -> SeqMap sym a)
-> (Integer -> SEval sym a) -> SeqMap sym a
forall a b. (a -> b) -> a -> b
$ \Integer
i ->
       case Integer -> Integer -> Maybe Integer
reindex Integer
i Integer
shft of
         Maybe Integer
Nothing -> SEval sym a
zro
         Just Integer
i' -> SeqMap sym a -> Integer -> SEval sym a
forall sym a. Backend sym => SeqMap sym a -> Integer -> SEval sym a
lookupSeqMap SeqMap sym a
vs Integer
i'


data IndexSegment sym
  = BitIndexSegment (SBit sym)
  | WordIndexSegment (SWord sym)

{-# SPECIALIZE
  barrelShifter ::
  Concrete ->
  (SBit Concrete -> a -> a -> SEval Concrete a) ->
  (SeqMap Concrete a -> Integer -> SEval Concrete (SeqMap Concrete a)) ->
  Nat' ->
  SeqMap Concrete a ->
  Integer ->
  [IndexSegment Concrete] ->
  SEval Concrete (SeqMap Concrete a)
 #-}
barrelShifter :: Backend sym =>
  sym ->
  (SBit sym -> a -> a -> SEval sym a)
     {- ^ if/then/else operation of values -} ->
  (SeqMap sym a -> Integer -> SEval sym (SeqMap sym a))
     {- ^ concrete shifting operation -} ->
  Nat' {- ^ Size of the map being shifted -} ->
  SeqMap sym a {- ^ initial value -} ->
  Integer {- Number of bits in shift amount -} ->
  [IndexSegment sym]  {- ^ segments of the shift amount, in big-endian order -} ->
  SEval sym (SeqMap sym a)
barrelShifter :: sym
-> (SBit sym -> a -> a -> SEval sym a)
-> (SeqMap sym a -> Integer -> SEval sym (SeqMap sym a))
-> Nat'
-> SeqMap sym a
-> Integer
-> [IndexSegment sym]
-> SEval sym (SeqMap sym a)
barrelShifter sym
sym SBit sym -> a -> a -> SEval sym a
mux SeqMap sym a -> Integer -> SEval sym (SeqMap sym a)
shift_op Nat'
sz SeqMap sym a
x0 Integer
n0 [IndexSegment sym]
bs0
  | Integer
n0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int) =
      IO (SeqMap sym a) -> SEval sym (SeqMap sym a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Unsupported -> IO (SeqMap sym a)
forall a e. Exception e => e -> a
X.throw (String -> Unsupported
UnsupportedSymbolicOp (String
"Barrel shifter with too many bits in shift amount: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
n0)))
  | Bool
otherwise = SeqMap sym a
-> Int -> [IndexSegment sym] -> SEval sym (SeqMap sym a)
go SeqMap sym a
x0 (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
n0) [IndexSegment sym]
bs0

  where
  go :: SeqMap sym a
-> Int -> [IndexSegment sym] -> SEval sym (SeqMap sym a)
go SeqMap sym a
x !Int
_n [] = SeqMap sym a -> SEval sym (SeqMap sym a)
forall (m :: * -> *) a. Monad m => a -> m a
return SeqMap sym a
x

  go SeqMap sym a
x !Int
n (WordIndexSegment SWord sym
w:[IndexSegment sym]
bs) =
    let n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Integer -> Int
forall a. Num a => Integer -> a
fromInteger (sym -> SWord sym -> Integer
forall sym. Backend sym => sym -> SWord sym -> Integer
wordLen sym
sym SWord sym
w) in
    case sym -> SWord sym -> Maybe (Integer, Integer)
forall sym.
Backend sym =>
sym -> SWord sym -> Maybe (Integer, Integer)
wordAsLit sym
sym SWord sym
w of
      Just (Integer
_,Integer
0) -> SeqMap sym a
-> Int -> [IndexSegment sym] -> SEval sym (SeqMap sym a)
go SeqMap sym a
x Int
n' [IndexSegment sym]
bs
      Just (Integer
_,Integer
j) ->
        do SeqMap sym a
x_shft <- SeqMap sym a -> Integer -> SEval sym (SeqMap sym a)
shift_op SeqMap sym a
x (Integer
j Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Bits a => Int -> a
bit Int
n')
           SeqMap sym a
-> Int -> [IndexSegment sym] -> SEval sym (SeqMap sym a)
go SeqMap sym a
x_shft Int
n' [IndexSegment sym]
bs
      Maybe (Integer, Integer)
Nothing ->
        do [SBit sym]
wbs <- sym -> SWord sym -> SEval sym [SBit sym]
forall sym. Backend sym => sym -> SWord sym -> SEval sym [SBit sym]
unpackWord sym
sym SWord sym
w
           SeqMap sym a
-> Int -> [IndexSegment sym] -> SEval sym (SeqMap sym a)
go SeqMap sym a
x Int
n ((SBit sym -> IndexSegment sym) -> [SBit sym] -> [IndexSegment sym]
forall a b. (a -> b) -> [a] -> [b]
map SBit sym -> IndexSegment sym
forall sym. SBit sym -> IndexSegment sym
BitIndexSegment [SBit sym]
wbs [IndexSegment sym] -> [IndexSegment sym] -> [IndexSegment sym]
forall a. [a] -> [a] -> [a]
++ [IndexSegment sym]
bs)

  go SeqMap sym a
x !Int
n (BitIndexSegment SBit sym
b:[IndexSegment sym]
bs) =
    let n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 in
    case sym -> SBit sym -> Maybe Bool
forall sym. Backend sym => sym -> SBit sym -> Maybe Bool
bitAsLit sym
sym SBit sym
b of
      Just Bool
False -> SeqMap sym a
-> Int -> [IndexSegment sym] -> SEval sym (SeqMap sym a)
go SeqMap sym a
x Int
n' [IndexSegment sym]
bs
      Just Bool
True ->
        do SeqMap sym a
x_shft <- SeqMap sym a -> Integer -> SEval sym (SeqMap sym a)
shift_op SeqMap sym a
x (Int -> Integer
forall a. Bits a => Int -> a
bit Int
n')
           SeqMap sym a
-> Int -> [IndexSegment sym] -> SEval sym (SeqMap sym a)
go SeqMap sym a
x_shft Int
n' [IndexSegment sym]
bs
      Maybe Bool
Nothing ->
        do SeqMap sym a
x_shft <- SeqMap sym a -> Integer -> SEval sym (SeqMap sym a)
shift_op SeqMap sym a
x (Int -> Integer
forall a. Bits a => Int -> a
bit Int
n')
           SeqMap sym a
x' <- sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
forall sym a.
Backend sym =>
sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
memoMap sym
sym Nat'
sz (sym
-> (SBit sym -> a -> a -> SEval sym a)
-> SBit sym
-> SeqMap sym a
-> SeqMap sym a
-> SeqMap sym a
forall sym a.
Backend sym =>
sym
-> (SBit sym -> a -> a -> SEval sym a)
-> SBit sym
-> SeqMap sym a
-> SeqMap sym a
-> SeqMap sym a
mergeSeqMap sym
sym SBit sym -> a -> a -> SEval sym a
mux SBit sym
b SeqMap sym a
x_shft SeqMap sym a
x)
           SeqMap sym a
-> Int -> [IndexSegment sym] -> SEval sym (SeqMap sym a)
go SeqMap sym a
x' Int
n' [IndexSegment sym]
bs