content: {- |
  Module           : Lang.JVM.Codebase
  Description      : Java Codebase manipulation
  License          : BSD3
  Stability        : stable
  Point-of-contact : jhendrix
  -}
  
  {-# OPTIONS_GHC -fno-warn-name-shadowing #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE DoAndIfThenElse #-}
  module Lang.JVM.Codebase
    ( Codebase
    , getClasses
    , isStrictSuper
    , isSubtype
    , loadCodebase
    , locateField
    , tryLookupClass
    , lookupClass
    , findStaticMethodsByRef
    , findVirtualMethodsByRef
    , supers
    , module Language.JVM.Parser
    ) where
  
  import Control.Monad
  import Control.Monad.IO.Class
  import Control.Monad.Trans.Maybe (MaybeT(..))
  import Crypto.Hash.SHA1 (hash)
  import qualified Data.ByteString.Base16 as B16
  import qualified Data.ByteString.Char8 as B
  import Data.Foldable (Foldable(..))
  import qualified Data.Map as M
  import Data.IORef
  import Data.List (isPrefixOf)
  import Data.Maybe
  import System.Directory
  import System.FilePath (pathSeparator, takeDirectory, (<.>), (</>))
  
  import Lang.JVM.JavaTools
  import Lang.JVM.ProcessUtils
  
  import Language.JVM.Common
  import Language.JVM.JarReader
  import Language.JVM.Parser
  
  {- Note [Loading classes from JIMAGE files]
  
  The Java standard library is packaged differently depending on which version of
  the JDK is being used:
  
  * JDK 8 and earlier ship the standard library in <java.home>/lib/rt.jar, so
    loading the standard library's .class files is simply a matter of parsing
    rt.jar. The `jvm-parser` library is used to parse rt.jar like any other JAR
    file.
  * JDK 9 and later abandon rt.jar and instead distribute the standard library as
    a JIMAGE file. (For more on the JIMAGE file format and how it compares to
    JAR, see https://stackoverflow.com/a/64202720.) In particular, the standard
    library is now located in a JIMAGE file at <java.home>/lib/modules.
  
  We need access to the standard library's .class files in order to verify even
  the simplest Java programs, so figuring out how to load classes from JIMAGE
  files is a requirement for using JDK 9 or later with crucible-jvm. Unlike JAR,
  however, the JIMAGE file format is internal to the JDK and subject to change in
  the future. As a result, parsing JIMAGE files directly, like `jvm-parser` does
  with JAR files, is likely not a good idea.
  
  If parsing JIMAGE files directly isn't a viable option, what can be done? As
  far as I can tell, there are only two reasonable ways of extracting .class
  files from JIMAGE files (outside of using Java itself):
  
  * Use libjimage, a C++ library bundled with the JDK, to load classes. This is
    what the JDK itself uses internally, and as a result, it's reasonably space-
    and time-efficient to load a single class from a JIMAGE file. Unfortunately,
    this is a non-starter for crucible-jvm due to potential linking issues that
    would arise when combining GHC with C++ code.
    See https://github.com/GaloisInc/crucible/issues/621.
  
  * Use the jimage utility, shipped alongside java in <java.home>/bin, to extract
    the contents of JIMAGE files. The `jimage extract` command can either be used
    to extract everything in a JIMAGE file all at once, or it can be used to
    extract one class at a time, like in this example:
  
      jimage extract --include regex:.*/java/lang/Class.class --dir <output-dir> <jimage-file>
  
  Since libjimage is an option, that only leaves jimage. Unfortunately, jimage is
  not without its flaws:
  
  * jimage is /slow/. It takes about about 1.5 seconds for jimage to extract the
    standard `modules` file on a reasonably performant machine. Having an extra
    1.5 seconds of lag every time that crucible-jvm is invoked doesn't sound
    appealing.
  
  * What's worse, jimage is still unreasonably slow even if you only use it to
    extract a single class from a JIMAGE file. It takes about 0.25 seconds to
    extract java/lang/Class.class from the standard `modules`, which leads me to
    suspect that jimage is doing a linear scan through every file in `modules`.
    In any case, this is about 14× slower than loading a class file from a JAR,
    and if crucible-jvm spends 0.25 seconds every time it loads any class from
    `modules`, this would add up very quickly.
  
  What can we do about jimage's performance issues? One appealing alternative is
  to use jlink, yet another utility that the JDK ships. If you give jlink a list
  of Java modules, it will produce a minimal runtime environment that includes a
  JIMAGE file that only contains the classes needed to support those modules, but
  nothing else. In theory, this could cut down on the time it takes to extract
  classes from a JIMAGE file, as the standard library's `modules` file is ~136 MB
  when compressed (~493 MB when uncompressed), and most Java classes will only
  use a small portion of that.
  
  Unfortunately, jlink is even slower than jimage. It takes jlink about 3 seconds
  to create a runtime environment that only contains the `java.base` module,
  which cancels out any time saved when extracting the JIMAGE file afterwards.
  
  Having exhausted all other possibilities, we have concluded that we simply
  can't make JIMAGE extraction fast in the general case. We /can/, however,
  amortize the cost of using JIMAGE files by caching the extracted contents.
  This is exactly what crucible-jvm does. Any time that crucible-jvm tries to
  load a class (e.g., java/lang/Class.class), it first checks to see if it is in
  a known cache directory
  (e.g., ~/.cache/crucible-jvm/<hashed-path-to-modules-file>/java.base/java/lang/Class.class),
  and if so, loads it like it would any other .class on the classpath. If not,
  it will then consult the JAR path and the classpath. If it is not located by
  that point, then crucible-jvm will attempt to extract it from the JIMAGE file,
  caching the result for subsequent lookups. Importantly, we perform the JIMAGE
  extraction step last, as it is far more time-consuming that trying to load
  classes from the JAR path or classpath.
  
  There's still possibly some room for further optimization here. As mentioned
  above, extracting one class at a time from a JIMAGE file takes about 0.25
  seconds, while extracting the entire JIMAGE file takes about 1.5 seconds.
  Depending on how many classes a program uses, it may take less time overall
  to just extract-and-cache everything from a JIMAGE file at once. The downside,
  however, is that the cache would require much more space. A typical JDK
  `modules` file takes up ~493 MB of classes when extracted, but most simple Java
  programs will only use <1 MB of these classes. As a result, we opt for the
  one-class-at-a-time approach, which seems better tuned for the sorts of
  Java programs that crucible-jvm verifies in common cases.
  -}
  
  -- | Collection of classes loaded by JVM.
  data CodebaseState = CodebaseState {
      classMap    :: M.Map ClassName Class
    , subclassMap :: M.Map ClassName [Class]
    -- ^ Maps class names to the list of classes that are direct subclasses, and
    -- interfaces to list of classes that directly implement them.
    }
  
  -- | Contains the path to the @jimage@ tool, the standard @modules@ JIMAGE
  -- file, and the path where the extracted contents of @modules@ are cached.
  --
  -- This is only used when crucible-jvm uses JDK 9 or later.
  -- See Note [Loading classes from JIMAGE files].
  data JimagePaths = JimagePaths
    { jimagePath               :: FilePath
    , standardModulesFilePath  :: FilePath
    , standardModulesCachePath :: FilePath
    }
  
  data Codebase = Codebase
    { jarReader  :: JarReader
    -- ^ Maps class names to lazily loaded classes in JARs
    , classPaths :: [FilePath]
    , jimagePaths :: Maybe JimagePaths
    -- ^ The path to the @jimage@ tool and the paths it works on.
    , stateRef   :: IORef CodebaseState
    }
  
  instance Show Codebase where
    show _ = "Codebase XXXXXX"
  
  -- | Loads Java classes directly beneath given path.  Also loads jar indices for
  -- lazy class loading.
  loadCodebase :: [FilePath] -> [FilePath] -> [FilePath] -> IO Codebase
  loadCodebase jarFiles classPaths javaBinDirs = do
    -- REVISIT: Currently, classes found in the classpath shadow those in the
    -- jars.  Pretty sure the -classpath argument to the vanilla jvm allows
    -- mixture of jars and directories, and resolves names in the order in which
    -- those elements are encountered.  We probably want to do the same thing (and
    -- be able to just provide one argument to loadCodebase), but not for the
    -- beta. [js 04 Nov 2010]
    --
    -- UPDATE: docs on class resolution:
    -- http://docs.oracle.com/javase/8/docs/technotes/tools/findingclasses.html.
    -- I don't see any mention of resolution order in case of name
    -- conflicts. This Stack Overflow answer claims it's complicated in
    -- general, but that it will be the first encountered, searching
    -- left to right in the class path:
    -- http://stackoverflow.com/a/9757708/470844.
    --
    -- If we later want to make the resolution order be the
    -- left-to-right-in-classpath order, then we can e.g. implement a
    -- 'ClassPathLoader' which includes sequence of 'JarLoader' and
    -- 'DirLoader' objects, which embed maps from class names to 'Class'
    -- objects. A
    --
    --   getClass :: String -> IO (Maybe Class)
    --
    -- interface would be sufficient. If we want to avoid repeating the
    -- lookup in many maps -- one for each classpath component -- we can
    -- merge the maps as in the current 'JarReader' type, but I doubt
    -- this would ever matter, performance wise.
    jars          <- newJarReader jarFiles
    mbJimagePaths <- findJimagePaths javaBinDirs
    let cb = CodebaseState { classMap = M.empty, subclassMap = M.empty }
    cbRef <- newIORef cb
    pure $ Codebase
      { jarReader   = jars
      , classPaths  = classPaths
      , jimagePaths = mbJimagePaths
      , stateRef    = cbRef
      }
  
  -- | Register a class with the given codebase
  addClass :: Class -> CodebaseState -> CodebaseState
  addClass cl (CodebaseState cMap scMap) =
    CodebaseState (M.insert (className cl) cl cMap)
                  (foldr addToSuperclass scMap
                     (maybeToList (superClass cl)++classInterfaces cl))
    where addToSuperclass super m =
            M.alter (\subclasses -> case subclasses of
                                      Just list -> Just (cl : list)
                                      Nothing -> Just [cl])
                    super
                    m
  
  -- | Returns class with given name in codebase or returns nothing if no class with
  -- that name can be found.
  tryLookupClass :: Codebase -> ClassName -> IO (Maybe Class)
  tryLookupClass (Codebase{jarReader, classPaths, jimagePaths, stateRef}) clNm = do
    cb <- readIORef stateRef
    case M.lookup clNm (classMap cb) of
      Just cl -> return (Just cl)
      Nothing -> do
        -- Here we bias our search to JARs before classpath directories,
        -- as mentioned above in 'loadCodebase'.
        let mcls = [ loadClassFromStandardModulesCache clNm jimagePaths
                   , loadClassFromJar clNm jarReader ] ++
                   map (loadClassFromDir clNm) classPaths ++
                   [ loadClassFromStandardModulesFile clNm jimagePaths ]
        mcl <- foldl1 firstSuccess mcls
        case mcl of
          Just cl -> do
            writeIORef stateRef $! addClass cl cb
            return $ Just cl
          Nothing -> return Nothing
  
  -- Search for a .class file in the cache of files extracted from the `modules`
  -- JIMAGE file. If it is not found, this does /not/ extract it from a JIMAGE
  -- file—see loadClassFromStandardModulesFile.
  loadClassFromStandardModulesCache :: ClassName -> Maybe JimagePaths -> IO (Maybe Class)
  loadClassFromStandardModulesCache clNm mbJimagePaths =
    case mbJimagePaths of
      Nothing -> pure Nothing
      Just JimagePaths{standardModulesCachePath} ->
        loadClassFromParentDir clNm standardModulesCachePath
  
  -- Attempt to extract a .class file from the `modules` JIMAGE file, cache it
  -- for subsequent lookups, and load the resulting .class.
  loadClassFromStandardModulesFile :: ClassName -> Maybe JimagePaths -> IO (Maybe Class)
  loadClassFromStandardModulesFile clNm mbJimagePaths =
    case mbJimagePaths of
      Nothing -> pure Nothing
      Just JimagePaths{jimagePath, standardModulesFilePath, standardModulesCachePath} -> do
        _ <- readProcessExitIfFailure jimagePath
               [ "extract"
               , "--include", "regex:.*/" ++ escapeDollars (unClassName clNm) ++ ".class"
               , "--dir", standardModulesCachePath
               , standardModulesFilePath
               ]
        -- After extracting a single class from a JIMAGE file, why do we then
        -- have to search every directory under the cached path again? It's
        -- because `jimage extract` doesn't tell you the path of the thing it
        -- just extracted. Sigh. We can't reasonably infer what this path is
        -- either, as it will start with a module name, which isn't something
        -- that is contained in a ClassName.
        loadClassFromParentDir clNm standardModulesCachePath
    where
      -- The '$' character needs to be escaped in the regex that
      -- `jimage extract` uses.
      escapeDollars :: String -> String
      escapeDollars []       = []
      escapeDollars ('$':cs) = '\\':'$':escapeDollars cs
      escapeDollars (c:cs)   = c:escapeDollars cs
  
  -- | Attempt to load a class by searching under directory @dir@, which
  -- is assumed to be a classpath component. If class @C1. ... .Cn@ is
  -- available under @dir@, then it must be located at
  -- @dir/C1/.../Cn.class@.
  -- http://docs.oracle.com/javase/8/docs/technotes/tools/findingclasses.html#userclass
  loadClassFromDir :: ClassName -> FilePath -> IO (Maybe Class)
  loadClassFromDir clNm dir = do
    exists <- doesFileExist file
    if exists
    then Just <$> loadClass file
    else return Nothing
    where
      file = dir </> classNameToClassFilePath clNm
      -- | Turn a @com/example/Class@-style classname into a
      --  @"com" </> "example" </> "Class.class"@-style platform dependent
      --  relative class-file path.
      --
      -- TODO: move this to 'jvm-parser.git:Language.JVM.Common'?
      classNameToClassFilePath :: ClassName -> FilePath
      classNameToClassFilePath clNm =
        map (\c -> if c == '/' then pathSeparator else c) (unClassName clNm) <.> "class"
  
  -- Given @parentDir@, search through its subdirectories to find a class,
  -- loading it if it is found.
  loadClassFromParentDir :: ClassName -> FilePath -> IO (Maybe Class)
  loadClassFromParentDir clNm parentDir = do
    childDirBaseNames <- listDirectory parentDir
    let childDirs = map (parentDir </>) childDirBaseNames
    foldl' firstSuccess (pure Nothing) $ map (loadClassFromDir clNm) childDirs
  
  -- | Attempt to find an executable named @jimage@, either in the directories
  -- supplied as arguments or on the @PATH@. If such an executable can be found,
  -- then @Just@ a @JimagePaths@ is returned. Otherwise, @Nothing@ is returned.
  --
  -- This is only used when crucible-jvm uses JDK 9 or later.
  -- See Note [Loading classes from JIMAGE files].
  findJimagePaths :: [FilePath] -> IO (Maybe JimagePaths)
  findJimagePaths javaBinDirs = runMaybeT $ do
    javaPath <- MaybeT $ findJavaIn javaBinDirs
    javaPathCanonical <- liftIO $ canonicalizePath javaPath -- Remove any symlinks
    jimagePath <- MaybeT $ findJimageIn [takeDirectory javaPathCanonical]
    javaHome <- MaybeT $ findJavaProperty javaPathCanonical "java.home"
    MaybeT $ do
      let standardModulesPath = javaHome </> "lib" </> "modules"
      xdgCacheDir <- getXdgDirectory XdgCache "crucible-jvm"
      let standardModulesHash = B16.encode $ hash $ B.pack standardModulesPath
          standardModulesCache = xdgCacheDir </> B.unpack standardModulesHash
      createDirectoryIfMissing True standardModulesCache
      pure $ Just $ JimagePaths
        { jimagePath               = jimagePath
        , standardModulesFilePath  = standardModulesPath
        , standardModulesCachePath = standardModulesCache
        }
  
  -- | Returns class with given name in codebase or raises error if no class with
  -- that name can be found.
  --
  -- The components of class name @clNm@ should be slash separated, not
  -- dot separated. E.g. the class @com.example.Class@ should be
  -- @com/example/class@.
  lookupClass :: Codebase -> ClassName -> IO Class
  lookupClass cb clNm = do
    maybeCl <- tryLookupClass cb clNm
    case maybeCl of
      Just cl -> return cl
      Nothing -> error $ errorMsg
    where
      dotNm = slashesToDots (unClassName clNm)
      isStandardLibClass = "java.lang" `isPrefixOf` dotNm
      errorMsg = unlines $
        if isStandardLibClass then
          [ "Cannot find class " ++ dotNm ++ " in codebase."
          , ""
          , "You probably forgot to specify the location of the"
          , "Java standard libraries JAR using the '-j' flag to saw or jss. The"
          , " standard libraries JAR is called 'classes.jar' on OS X systems and"
          , "'rt.jar' on Windows and Linux systems. Its location can be found by"
          , "running 'java -verbose 2>&1 | grep Opened', assuming you're using"
          , "a Sun Java."
          ]
        else
          [ "Cannot find class " ++ dotNm ++ " in codebase."
          , ""
          , "You can specify the location of classes you depend on using"
          , "the '-c' flag to specify non-jar classpaths and the '-j' flag"
          , "to specify the location of JAR files."
          ]
  
  getClasses :: Codebase -> IO [Class]
  getClasses (Codebase{stateRef}) = do
    cb <- readIORef stateRef
    return . M.elems . classMap $ cb
  
  -- | Adjusts the given field id to specify as its class the class in the
  -- superclass hierarchy that declares it
  locateField :: Codebase -> FieldId -> IO FieldId
  locateField cb fldId = do
    owner <- findFieldOwner fldId
    return $ fldId { fieldIdClass = className owner}
    where
      -- Walk an inheritance hierarchy to determine the the class that declares a
      -- given field (i.e., the "field owner")
      findFieldOwner :: FieldId -> IO Class
      findFieldOwner FieldId{fieldIdName = fldNm, fieldIdClass = clNm } = do
        sups <- supers cb =<< lookupClass cb clNm
        case filter hasField sups of
          -- In theory, this should be unreachable.
          [] -> error $ "Unable to find field '" ++ fldNm
                      ++ "' in superclass hierarchy of class " ++ unClassName clNm
          (cl' : _) -> return cl'
        where
          hasField cl = fldNm `elem` map fieldName accessibleFields
            where accessibleFields
                    | className cl == clNm = classFields cl
                    | otherwise =
                        filter ((/=) Private . fieldVisibility) (classFields cl)
  
  -- | (isStrictSuper cb name class) returns true if name is a (strict) superclass
  -- of class in cb.
  isStrictSuper :: Codebase -> ClassName -> Class -> IO Bool
  isStrictSuper cb name cl = do
    case superClass cl of
      Just super
        | name == super -> return True
        | otherwise -> isStrictSuper cb name =<< lookupClass cb super
      Nothing -> return False
  
  -- | Returns true if subclass is a subtype of superclass in codebase.
  isSubtype :: Codebase -> Type -> Type -> IO Bool
  isSubtype _ sub super | sub == super           = return True
  isSubtype cb (ArrayType sub) (ArrayType super)  = isSubtype cb sub super
  isSubtype _ (ArrayType _sub) (ClassType super) =
    return $ super == "java/lang/Object"
             || super == "java/lang/Cloneable"
             || super == "java/io/Serializable"
  isSubtype cb (ClassType subName) super@(ClassType _) = do
   subclass <- lookupClass cb subName
   -- Check if sub is a subclass of super
   b <- case (superClass subclass) of
          Just superName -> isSubtype cb (ClassType superName) super
          Nothing        -> return False
   -- Check if super is an interface that sub implements
   b' <- or <$> mapM (\i -> isSubtype cb (ClassType i) super)
                     (classInterfaces subclass)
   return $ b || b'
  isSubtype _ _sub _super = return False
  
  -- | Finds all classes that implement a given method. List is ordered so that
  -- subclasses appear before their base class.
  findVirtualMethodsByRef :: Codebase
                          -> ClassName -- ^ Name of class given in code for this class.
                          -> MethodKey -- ^ Method to identify.
                          -> ClassName -- ^ Concrete class for this object.
                          -> IO [ClassName]
  findVirtualMethodsByRef cb name key instTyNm = do
    cl <- lookupClass cb name
    sups    <- drop 1 <$> supers cb cl
    revSubs <- reverse <$> cb `subs` cl
    map className <$> filterM isMatch (revSubs ++ sups)
   where
      isMatch cl = case cl `lookupMethod` key of
        Nothing -> return False
        Just _  -> do
          b   <- isStrictSuper cb (className cl) =<< lookupClass cb instTyNm
          return $ className cl == instTyNm || b
  
  -- | Finds all classes that implement a given static method. List is
  -- ordered so that subclasses appear before their base class.
  findStaticMethodsByRef :: Codebase
                         -> ClassName
                         -- ^ Name of class given in the code for this call
                         -> MethodKey
                         -- ^ Method to call
                         -> IO [ClassName]
  findStaticMethodsByRef cb name key = do
    cl <- lookupClass cb name
    sups <- supers cb cl
    let isMatch clName = isJust (clName `lookupMethod` key)
    return . map className . filter isMatch $ sups
  
  -- | Produces the superclass hierarchy of the given class. Ordered from subclass
  -- to base class, starting with the given class.
  supers :: Codebase -> Class -> IO [Class]
  supers cb cl = do
    starClosureM (maybe (return []) (fmap (:[]) . lookupClass cb) . superClass) cl
  
  -- | Produces the subclass hierarchy of the given class.  Ordered
  -- from base class to subclass, starting with the given class.
  subs :: Codebase -> Class -> IO [Class]
  subs (Codebase{stateRef}) cl = do
    cb <- readIORef stateRef
    return $ starClosure (fromMaybe [] . (`M.lookup` subclassMap cb) . className) cl
  
  
  --------------------------------------------------------------------------------
  -- Misc
  
  -- (starClosure f a) returns the list of values obtainable by 0 or more
  -- applications of f.
  starClosure :: (a -> [a]) -> a -> [a]
  starClosure fn a = a : concatMap (starClosure fn) (fn a)
  
  -- The monadic variant of starClosure
  starClosureM :: Monad m => (a -> m [a]) -> a -> m [a]
  starClosureM fn a =
    return ((a:) . concat) `ap` (mapM (starClosureM fn) =<< fn a)
  
  -- | Combine two @IO (Maybe a)@ computations lazily, choosing the
  -- first to succeed (i.e. return 'Just').
  firstSuccess :: IO (Maybe a) -> IO (Maybe a) -> IO (Maybe a)
  -- This seems like it would be a common pattern, although I can't
  -- find it in the standard libraries.
  firstSuccess ima1 ima2 = do
    ma1 <- ima1
    case ma1 of
      Nothing -> ima2
      Just _ -> return ma1
  