{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TemplateHaskell            #-}
{-# LANGUAGE TupleSections              #-}
{-# LANGUAGE ViewPatterns               #-}
{-|
Copyright   : (c) Christian Gram Kalhauge, 2018
License     : MIT
Maintainer  : kalhuage@cs.ucla.edu

This module contains the 'ConstantPool' data structure and multiple
other types, and classes.
-}
module Language.JVM.ConstantPool
  (
  -- * Constant Pool
  -- $ConstantPool
    ConstantPool (..)
  , access

  , growPool

  , poolCount
  , nextIndex
  , listConstants
  , fromConstants

  , empty
  , PoolAccessError (..)

  , Index
  ) where

import           Control.DeepSeq       (NFData)
import           Control.Monad.Except
import           Data.Binary
-- import           Debug.Trace
import           Data.Binary.Get
import           Data.Binary.Put
import           GHC.Generics          (Generic)

-- base
import Data.Monoid

-- containers
import qualified Data.IntMap           as IM

import           Language.JVM.Constant
import           Language.JVM.Stage
import           Language.JVM.TH


-- $ConstantPool
-- The 'ConstantPool' contains all the constants, and is accessible using the
-- Lookup methods.

-- | A ConstantPool is just an 'IntMap'. A 'IntMap' is used, because constants are
-- accessed using their byte-offset, and sometimes the offset depends on the constant
-- size. See 'constantSize'.
newtype ConstantPool r = ConstantPool
  { ConstantPool r -> IntMap (Constant r)
unConstantPool :: IM.IntMap (Constant r)
  }

instance Binary (ConstantPool Low) where
  get :: Get (ConstantPool Low)
get = do
    Index
len <- Index -> Index
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Index -> Index) -> Get Index -> Get Index
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Index
getWord16be
    [(Key, Constant Low)]
list <- Index -> Index -> Get [(Key, Constant Low)]
forall r a.
(Binary (Constant r), Num a) =>
Index -> Index -> Get [(a, Constant r)]
go Index
len Index
1
    ConstantPool Low -> Get (ConstantPool Low)
forall (m :: * -> *) a. Monad m => a -> m a
return (ConstantPool Low -> Get (ConstantPool Low))
-> (IntMap (Constant Low) -> ConstantPool Low)
-> IntMap (Constant Low)
-> Get (ConstantPool Low)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (Constant Low) -> ConstantPool Low
forall r. IntMap (Constant r) -> ConstantPool r
ConstantPool (IntMap (Constant Low) -> Get (ConstantPool Low))
-> IntMap (Constant Low) -> Get (ConstantPool Low)
forall a b. (a -> b) -> a -> b
$ [(Key, Constant Low)] -> IntMap (Constant Low)
forall a. [(Key, a)] -> IntMap a
IM.fromList [(Key, Constant Low)]
list
    where
      go :: Index -> Index -> Get [(a, Constant r)]
go Index
len Index
n | Index
len Index -> Index -> Bool
forall a. Ord a => a -> a -> Bool
> Index
n = do
        Constant r
constant <- Get (Constant r)
forall t. Binary t => Get t
get
        [(a, Constant r)]
rest <- Index -> Index -> Get [(a, Constant r)]
go Index
len (Index
n Index -> Index -> Index
forall a. Num a => a -> a -> a
+ Constant r -> Index
forall r. Constant r -> Index
constantSize Constant r
constant)
        [(a, Constant r)] -> Get [(a, Constant r)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(a, Constant r)] -> Get [(a, Constant r)])
-> [(a, Constant r)] -> Get [(a, Constant r)]
forall a b. (a -> b) -> a -> b
$ (Index -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Index
n, Constant r
constant) (a, Constant r) -> [(a, Constant r)] -> [(a, Constant r)]
forall a. a -> [a] -> [a]
: [(a, Constant r)]
rest
      go Index
_ Index
_ = [(a, Constant r)] -> Get [(a, Constant r)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
  put :: ConstantPool Low -> Put
put (ConstantPool IntMap (Constant Low)
p) = do
    case IntMap (Constant Low)
-> Maybe ((Key, Constant Low), IntMap (Constant Low))
forall a. IntMap a -> Maybe ((Key, a), IntMap a)
IM.maxViewWithKey IntMap (Constant Low)
p of
      Just ((Key
key, Constant Low
e), IntMap (Constant Low)
_) -> do
        Index -> Put
putWord16be (Key -> Index
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
key Index -> Index -> Index
forall a. Num a => a -> a -> a
+ Constant Low -> Index
forall r. Constant r -> Index
constantSize Constant Low
e)
        [(Key, Constant Low)] -> ((Key, Constant Low) -> Put) -> Put
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (IntMap (Constant Low) -> [(Key, Constant Low)]
forall a. IntMap a -> [(Key, a)]
IM.toAscList IntMap (Constant Low)
p) (Constant Low -> Put
forall t. Binary t => t -> Put
put (Constant Low -> Put)
-> ((Key, Constant Low) -> Constant Low)
-> (Key, Constant Low)
-> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Key, Constant Low) -> Constant Low
forall a b. (a, b) -> b
snd)
      Maybe ((Key, Constant Low), IntMap (Constant Low))
Nothing -> do
        Int16 -> Put
putInt16be Int16
0

-- | A pool access error
data PoolAccessError = PoolAccessError
  { PoolAccessError -> Index
paErrorRef :: !Word16
  , PoolAccessError -> String
paErrorMsg :: String
  } deriving (Key -> PoolAccessError -> ShowS
[PoolAccessError] -> ShowS
PoolAccessError -> String
(Key -> PoolAccessError -> ShowS)
-> (PoolAccessError -> String)
-> ([PoolAccessError] -> ShowS)
-> Show PoolAccessError
forall a.
(Key -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PoolAccessError] -> ShowS
$cshowList :: [PoolAccessError] -> ShowS
show :: PoolAccessError -> String
$cshow :: PoolAccessError -> String
showsPrec :: Key -> PoolAccessError -> ShowS
$cshowsPrec :: Key -> PoolAccessError -> ShowS
Show, PoolAccessError -> PoolAccessError -> Bool
(PoolAccessError -> PoolAccessError -> Bool)
-> (PoolAccessError -> PoolAccessError -> Bool)
-> Eq PoolAccessError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PoolAccessError -> PoolAccessError -> Bool
$c/= :: PoolAccessError -> PoolAccessError -> Bool
== :: PoolAccessError -> PoolAccessError -> Bool
$c== :: PoolAccessError -> PoolAccessError -> Bool
Eq, (forall x. PoolAccessError -> Rep PoolAccessError x)
-> (forall x. Rep PoolAccessError x -> PoolAccessError)
-> Generic PoolAccessError
forall x. Rep PoolAccessError x -> PoolAccessError
forall x. PoolAccessError -> Rep PoolAccessError x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep PoolAccessError x -> PoolAccessError
$cfrom :: forall x. PoolAccessError -> Rep PoolAccessError x
Generic)

instance NFData PoolAccessError

-- | Creates an empty constant pool
empty :: ConstantPool r
empty :: ConstantPool r
empty = IntMap (Constant r) -> ConstantPool r
forall r. IntMap (Constant r) -> ConstantPool r
ConstantPool (IntMap (Constant r)
forall a. IntMap a
IM.empty)

-- | Access a constant in the constant pool
access :: Index -> ConstantPool r -> Either PoolAccessError (Constant r)
access :: Index -> ConstantPool r -> Either PoolAccessError (Constant r)
access Index
ref (ConstantPool IntMap (Constant r)
cp) =
  case Key -> IntMap (Constant r) -> Maybe (Constant r)
forall a. Key -> IntMap a -> Maybe a
IM.lookup (Index -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Index
ref) IntMap (Constant r)
cp of
    Just Constant r
x  -> Constant r -> Either PoolAccessError (Constant r)
forall a b. b -> Either a b
Right Constant r
x
    Maybe (Constant r)
Nothing -> PoolAccessError -> Either PoolAccessError (Constant r)
forall a b. a -> Either a b
Left (PoolAccessError -> Either PoolAccessError (Constant r))
-> PoolAccessError -> Either PoolAccessError (Constant r)
forall a b. (a -> b) -> a -> b
$ Index -> String -> PoolAccessError
PoolAccessError Index
ref String
"No such element."

poolCount :: ConstantPool r -> Int
poolCount :: ConstantPool r -> Key
poolCount =
  IntMap (Constant r) -> Key
forall a. IntMap a -> Key
IM.size (IntMap (Constant r) -> Key)
-> (ConstantPool r -> IntMap (Constant r)) -> ConstantPool r -> Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConstantPool r -> IntMap (Constant r)
forall r. ConstantPool r -> IntMap (Constant r)
unConstantPool

listConstants :: ConstantPool r -> [(Index, Constant r)]
listConstants :: ConstantPool r -> [(Index, Constant r)]
listConstants =
  ((Key, Constant r) -> (Index, Constant r))
-> [(Key, Constant r)] -> [(Index, Constant r)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Key
i, Constant r
a) -> (Key -> Index
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
i, Constant r
a)) ([(Key, Constant r)] -> [(Index, Constant r)])
-> (ConstantPool r -> [(Key, Constant r)])
-> ConstantPool r
-> [(Index, Constant r)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (Constant r) -> [(Key, Constant r)]
forall a. IntMap a -> [(Key, a)]
IM.toList (IntMap (Constant r) -> [(Key, Constant r)])
-> (ConstantPool r -> IntMap (Constant r))
-> ConstantPool r
-> [(Key, Constant r)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConstantPool r -> IntMap (Constant r)
forall r. ConstantPool r -> IntMap (Constant r)
unConstantPool

nextIndex :: ConstantPool r -> Index
nextIndex :: ConstantPool r -> Index
nextIndex (ConstantPool IntMap (Constant r)
im) =
  Index -> Index
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Index -> Index) -> Index -> Index
forall a b. (a -> b) -> a -> b
$ case IntMap (Constant r) -> [(Key, Constant r)]
forall a. IntMap a -> [(Key, a)]
IM.toDescList IntMap (Constant r)
im of
    (Key
k, Constant r
a):[(Key, Constant r)]
_ -> Key -> Index
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
k Index -> Index -> Index
forall a. Num a => a -> a -> a
+ Constant r -> Index
forall r. Constant r -> Index
constantSize Constant r
a
    [(Key, Constant r)]
_        -> Index
1

fromConstants :: Foldable f => f (Constant r) -> ConstantPool r
fromConstants :: f (Constant r) -> ConstantPool r
fromConstants =
  (ConstantPool r -> Constant r -> ConstantPool r)
-> ConstantPool r -> f (Constant r) -> ConstantPool r
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\ConstantPool r
b Constant r
a -> (Index, ConstantPool r) -> ConstantPool r
forall a b. (a, b) -> b
snd ((Index, ConstantPool r) -> ConstantPool r)
-> (ConstantPool r -> (Index, ConstantPool r))
-> ConstantPool r
-> ConstantPool r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Constant r -> ConstantPool r -> (Index, ConstantPool r)
forall r. Constant r -> ConstantPool r -> (Index, ConstantPool r)
append Constant r
a (ConstantPool r -> ConstantPool r)
-> ConstantPool r -> ConstantPool r
forall a b. (a -> b) -> a -> b
$ ConstantPool r
b) ConstantPool r
forall r. ConstantPool r
empty

growPool ::
  forall b.
  (ConstantPool High -> Constant Low -> Either b (Constant High))
  -> ConstantPool Low
  -> (ConstantPool High, [(b, (Index, Constant Low))])
growPool :: (ConstantPool High -> Constant Low -> Either b (Constant High))
-> ConstantPool Low
-> (ConstantPool High, [(b, (Index, Constant Low))])
growPool ConstantPool High -> Constant Low -> Either b (Constant High)
f ConstantPool Low
reffed =
  IntMap (Constant High)
-> [(Index, Constant Low)]
-> (ConstantPool High, [(b, (Index, Constant Low))])
stage' IntMap (Constant High)
forall a. IntMap a
IM.empty (ConstantPool Low -> [(Index, Constant Low)]
forall r. ConstantPool r -> [(Index, Constant r)]
listConstants ConstantPool Low
reffed)
  where
    stage' :: IM.IntMap (Constant High) -> [(Index, Constant Low)] -> (ConstantPool High, [(b, (Index, Constant Low))])
    stage' :: IntMap (Constant High)
-> [(Index, Constant Low)]
-> (ConstantPool High, [(b, (Index, Constant Low))])
stage' IntMap (Constant High)
cp [(Index, Constant Low)]
mis =
      case ((Index, Constant Low)
 -> (IntMap (Constant High), Endo [(b, (Index, Constant Low))]))
-> [(Index, Constant Low)]
-> (IntMap (Constant High), Endo [(b, (Index, Constant Low))])
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (ConstantPool High
-> (Index, Constant Low)
-> (IntMap (Constant High), Endo [(b, (Index, Constant Low))])
forall a.
Integral a =>
ConstantPool High
-> (a, Constant Low)
-> (IntMap (Constant High), Endo [(b, (a, Constant Low))])
grow (IntMap (Constant High) -> ConstantPool High
forall r. IntMap (Constant r) -> ConstantPool r
ConstantPool IntMap (Constant High)
cp)) [(Index, Constant Low)]
mis of
        (IntMap (Constant High)
cp', (Endo [(b, (Index, Constant Low))]
 -> [(b, (Index, Constant Low))] -> [(b, (Index, Constant Low))])
-> [(b, (Index, Constant Low))]
-> Endo [(b, (Index, Constant Low))]
-> [(b, (Index, Constant Low))]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Endo [(b, (Index, Constant Low))]
-> [(b, (Index, Constant Low))] -> [(b, (Index, Constant Low))]
forall a. Endo a -> a -> a
appEndo [] -> [(b, (Index, Constant Low))]
mis')
          | IntMap (Constant High) -> Bool
forall a. IntMap a -> Bool
IM.null IntMap (Constant High)
cp' ->
              (IntMap (Constant High) -> ConstantPool High
forall r. IntMap (Constant r) -> ConstantPool r
ConstantPool IntMap (Constant High)
cp, [(b, (Index, Constant Low))]
mis')
          | Bool
otherwise ->
          IntMap (Constant High)
-> [(Index, Constant Low)]
-> (ConstantPool High, [(b, (Index, Constant Low))])
stage' (IntMap (Constant High)
cp IntMap (Constant High)
-> IntMap (Constant High) -> IntMap (Constant High)
forall a. IntMap a -> IntMap a -> IntMap a
`IM.union` IntMap (Constant High)
cp') ([(Index, Constant Low)]
 -> (ConstantPool High, [(b, (Index, Constant Low))]))
-> ([(b, (Index, Constant Low))] -> [(Index, Constant Low)])
-> [(b, (Index, Constant Low))]
-> (ConstantPool High, [(b, (Index, Constant Low))])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((b, (Index, Constant Low)) -> (Index, Constant Low))
-> [(b, (Index, Constant Low))] -> [(Index, Constant Low)]
forall a b. (a -> b) -> [a] -> [b]
map (b, (Index, Constant Low)) -> (Index, Constant Low)
forall a b. (a, b) -> b
snd ([(b, (Index, Constant Low))]
 -> (ConstantPool High, [(b, (Index, Constant Low))]))
-> [(b, (Index, Constant Low))]
-> (ConstantPool High, [(b, (Index, Constant Low))])
forall a b. (a -> b) -> a -> b
$ [(b, (Index, Constant Low))]
mis'

    grow :: ConstantPool High
-> (a, Constant Low)
-> (IntMap (Constant High), Endo [(b, (a, Constant Low))])
grow ConstantPool High
cp (a
k,Constant Low
a) =
      case ConstantPool High -> Constant Low -> Either b (Constant High)
f ConstantPool High
cp Constant Low
a of
        Right Constant High
c -> (Key -> Constant High -> IntMap (Constant High)
forall a. Key -> a -> IntMap a
IM.singleton (a -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
k) Constant High
c, Endo [(b, (a, Constant Low))]
forall a. Monoid a => a
mempty)
        Left b
b  -> (IntMap (Constant High)
forall a. IntMap a
IM.empty, ([(b, (a, Constant Low))] -> [(b, (a, Constant Low))])
-> Endo [(b, (a, Constant Low))]
forall a. (a -> a) -> Endo a
Endo ((b
b, (a
k,Constant Low
a))(b, (a, Constant Low))
-> [(b, (a, Constant Low))] -> [(b, (a, Constant Low))]
forall a. a -> [a] -> [a]
:) )


{-# INLINE growPool #-}

-- | Append a constant to the constant pool, and get the offset.
append :: Constant r -> ConstantPool r -> (Index, ConstantPool r)
append :: Constant r -> ConstantPool r -> (Index, ConstantPool r)
append Constant r
c cp :: ConstantPool r
cp@(ConstantPool IntMap (Constant r)
im) =
  (Index
i, IntMap (Constant r) -> ConstantPool r
forall r. IntMap (Constant r) -> ConstantPool r
ConstantPool (IntMap (Constant r) -> ConstantPool r)
-> IntMap (Constant r) -> ConstantPool r
forall a b. (a -> b) -> a -> b
$ Key -> Constant r -> IntMap (Constant r) -> IntMap (Constant r)
forall a. Key -> a -> IntMap a -> IntMap a
IM.insert (Index -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Index
i) Constant r
c IntMap (Constant r)
im)
  where
    i :: Index
i = ConstantPool r -> Index
forall r. ConstantPool r -> Index
nextIndex ConstantPool r
cp

$(deriveBase ''ConstantPool)