content: -- | This plugin generates Java bytecode from modules using the java
  -- QuasiQuoter and inserts it in a global bytecode table from where
  -- it is loaded at runtime.
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE ViewPatterns #-}
  module Language.Java.Inline.Plugin (plugin) where
  
  import Control.Monad (forM, when, zipWithM)
  import Control.Monad.Writer
  import qualified Data.ByteString as BS
  import Data.ByteString.Builder (Builder)
  import qualified Data.ByteString.Builder as Builder
  import Data.Char (chr, ord)
  import Data.List (find, intersperse, isSuffixOf)
  import Data.Maybe (fromMaybe)
  import Data.Monoid (Endo(..))
  import qualified Data.Text as Text
  import qualified Data.Text.Encoding as Text
  import Foreign.JNI.Types (JType(..))
  import GhcPlugins.Extras
  import qualified Language.Haskell.TH as TH
  import qualified Language.Haskell.TH.Syntax as TH
  import Language.Java.Inline.Internal.Magic
  import Language.Java.Inline.Internal.QQMarker.Names (getQQMarkers)
  import System.Directory (listDirectory)
  import System.FilePath ((</>), (<.>), takeDirectory)
  import System.IO (withFile, IOMode(WriteMode), hPutStrLn, stderr)
  import System.IO.Temp (withSystemTempDirectory)
  import System.Process (callProcess)
  import Prelude hiding ((<>))
  
  -- The 'java' quasiquoter produces annotations of type 'JavaImport', and it also
  -- inserts calls in the code to the function 'qqMarker'.
  --
  -- 'qqMarker' carries many bits of information that are useful in generating the
  -- QQ code.
  --
  -- This plugin first makes a pass to collect the 'qqMarker' calls in the module
  -- (collectQQMarkers).
  --
  -- Then it translates the Core Types to Java types (unliftJTypes).
  --
  -- Then it generates the java stubs from the information extracted from the
  -- occurrences of 'qqMarker' (buildJava).
  --
  -- Finally, it calls the java compiler to produce the bytecode and
  -- arranges to have it inserted in the bytecode table in constructor functions
  -- (cConstructor).
  
  plugin :: Plugin
  plugin = defaultPlugin
      { installCoreToDos = install
      , pluginRecompile = flagRecompile
      }
    where
      install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
      install args todo = do
        let passName = "inline-java"
        if inTodo passName todo then return todo
        else return (CoreDoPluginPass passName (qqPass args) : todo)
  
      inTodo :: String -> [CoreToDo] -> Bool
      inTodo name = \case
        CoreDoPluginPass n _ : xs -> n == name || inTodo name xs
        _ : xs -> inTodo name xs
        [] -> False
  
      qqPass :: [CommandLineOption] -> ModGuts -> CoreM ModGuts
      qqPass args guts = do
        getQQMarkers >>= \case
          -- If qqMarkers cannot be found we assume the module does not use
          -- inline-java.
          [] -> return guts
          qqMarkerNames -> do
            (binds, qqOccs) <- collectQQMarkers qqMarkerNames (mg_binds guts)
            let jimports =
                  GhcPlugins.Extras.getModuleAnnotations guts :: [JavaImport]
            dcs <- buildJava guts qqOccs jimports
                     >>= maybeDumpJava args
                     >>= buildBytecode args
            dflags <- getDynFlags
            let cstub =
                  cConstructor
                    dflags
                    (mg_module guts)
                    (text bctable_header $$ dotClasses dcs)
            return guts
              { mg_binds = binds
              , mg_foreign = appendStubC (mg_foreign guts) cstub
              }
  
      -- The contents of bctable.h
      --
      -- #include "bctable.h" wouldn't work when ghc is used from the
      -- command line without saying -package inline-java.
      bctable_header :: String
      bctable_header = $(do
          loc <- TH.location
          let parents = iterate takeDirectory (TH.loc_filename loc)
              depth = 5
              almostRoot = parents !! (depth - 1)
              root = parents !! depth
              f = root </> "cbits/bctable.h"
          when (root == almostRoot) $
            fail "root reached prematurely"
          TH.addDependentFile f
          TH.lift =<< TH.runIO (readFile f)
        )
      -- Dumps the java code to stderr or a file, depending on the set flags.
      maybeDumpJava :: [CommandLineOption] -> Builder -> CoreM Builder
      maybeDumpJava args b
        | elem "dump-java" args = do
          dflags <- getDynFlags
          if gopt Opt_DumpToFile dflags then do
            thisModule <- getModule
            let fname = moduleNameString (moduleName thisModule) ++ ".dump-java"
                path = maybe fname (</> fname) (dumpDir dflags)
            liftIO $ withFile path WriteMode $ \h -> Builder.hPutBuilder h b
          else liftIO $ do
            hPutStrLn stderr "=== inline-java (dump-java) BEGIN ==="
            Builder.hPutBuilder stderr b
            hPutStrLn stderr "=== inline-java (dump-java) END ==="
          return b
      maybeDumpJava _ b = return b
  
  -- | Produces a Java compilation unit from the quasiquotation occurrences and
  -- the java imports.
  --
  -- The compilation unit looks like:
  --
  -- > package io.tweag.inlinejava;
  -- > import java.util.*; // .hs:25
  -- >
  -- > public final class Inline__<pkgname>_<modname> {
  -- > public static java.lang.Object method_0()
  -- > { return  1 + 1 ; } // .hs:31
  -- > public static java.lang.Object inline__method_1()
  -- > { // .hs:34
  -- >              int x = 1; // .hs:34
  -- >              int y = 2; // .hs:34
  -- >              return x + y; // .hs:34
  -- >            } // .hs:34
  -- > public static java.lang.Object inline__method_2(final int $x)
  -- > { return  $x + 1 ; } // .hs:42
  -- > public static java.lang.Object inline__method_3()
  -- > { return  new Object() {} ; } // .hs:46
  -- > }
  --
  -- Where @inline_method_i@ is the method corresponding to the @ith@
  -- quasiquotation.
  buildJava :: ModGuts -> [QQOcc] -> [JavaImport] -> CoreM Builder
  buildJava guts qqOccs jimports = do
      let importsJava = mconcat
            [ mconcat [ "import ", Builder.stringUtf8 jimp
                      , "; // .hs:", Builder.integerDec n
                      , "\n"
                      ]
            | JavaImport jimp n <- jimports
            ]
      p_fam_env <- getPackageFamInstEnv
      let fam_envs = (p_fam_env, mg_fam_inst_env guts)
      methods <- forM qqOccs $ \QQOcc {..} -> do
        let Reduction _ normty = normaliseType fam_envs Nominal (expandTypeSynonyms qqOccResTy)
        jTypeNames <- findJTypeNames
        resty <- case toJavaType jTypeNames normty of
          Just resty -> return resty
          Nothing -> GhcPlugins.Extras.failWith $ hsep
            [ parens (text "line" <+> integer qqOccLineNumber) <> ":"
            , text "The result type of the quasiquotation"
            , quotes (ppr qqOccResTy)
            , text "is not sufficiently instantiated to infer a java type."
            ]
        let argnames = BS.split (fromIntegral $ fromEnum ',') qqOccAntiQs
        argtys <- zipWithM (getArg jTypeNames qqOccLineNumber)
                           argnames qqOccArgTys
        return $ mconcat
          [ "public static "
          , Builder.byteString resty
          , " "
          , Builder.byteString qqOccMName
          , "("
          , mconcat $ intersperse "," argtys
          , ") throws Throwable // .hs:"
          , Builder.integerDec qqOccLineNumber
          , "\n"
          , adjustInput qqOccLineNumber qqOccInput
          ]
      thisModule <- getModule
      let className = mangle thisModule
      return $ mconcat
        [ "package io.tweag.inlinejava;\n"
        , importsJava
        , "\n"
        , "public final class "
        , Builder.stringUtf8 className
        , " {\n"
        , mconcat methods
        , "}\n"
        ]
    where
      getArg :: JTypeNames -> Integer -> BS.ByteString -> Type -> CoreM Builder
      getArg jTypeNames _ name
             (expandTypeSynonyms -> toJavaType jTypeNames -> Just jtype) =
        return $ mconcat
          ["final ", Builder.byteString jtype, " $", Builder.byteString name]
      getArg _ lineN name t = GhcPlugins.Extras.failWith $ hsep
          [ parens (text "line" <+> integer lineN) <> ":"
          , quotes (ftext (mkFastStringByteString name) <+> "::" <+> ppr t)
          , text "is not sufficiently instantiated to infer a java type."
          ]
  
      adjustInput :: Integer -> BS.ByteString -> Builder
      adjustInput lineNumber bs =
        let txt = Text.strip $ Text.decodeUtf8 bs
            block = if Text.isPrefixOf "{" txt && Text.isSuffixOf "}" txt
              then Text.lines txt
              else Text.lines $ Text.concat ["{ return ", txt, "}" ]
  
         in Builder.byteString $ Text.encodeUtf8 $ Text.unlines
              [ Text.concat [ln, " // .hs:", Text.pack (show lineNumber)]
              | ln <- block
              ]
  
  -- | Produces a class name from a Module.
  mangle :: Module -> String
  mangle m = mangleClassName (unitIdString $ moduleUnitId m)
                             (moduleNameString (moduleName m))
  
  -- Call the java compiler and feeds it the given Java code in Builder form.
  buildBytecode :: [CommandLineOption] -> Builder -> CoreM [DotClass]
  buildBytecode args unit = do
      let javac = fromMaybe "javac" $ find ("javac" `isSuffixOf`) args
      m <- getModule
      liftIO $ withSystemTempDirectory "inlinejava" $ \dir -> do
        let src = dir </> mangle m <.> "java"
        withFile src WriteMode $ \h -> Builder.hPutBuilder h unit
        callProcess javac [src]
        -- A single compilation unit can produce multiple class files.
        classFiles <- filter (".class" `isSuffixOf`) <$> listDirectory dir
        forM classFiles $ \classFile -> do
          bcode <- BS.readFile (dir </> classFile)
          -- Strip the .class suffix.
          let klass = "io.tweag.inlinejava." ++ takeWhile (/= '.') classFile
          return $ DotClass klass bcode
  
  -- | The names of 'JType' data constructors
  data JTypeNames = JTypeNames
      { nameClass :: Maybe Name
      , nameIface :: Maybe Name
      , nameArray :: Maybe Name
      , nameGeneric :: Maybe Name
      , namePrim :: Maybe Name
      , nameVoid :: Maybe Name
      }
  
  -- | Produces the names of the data constructors of the 'JType'
  -- if they are used in the current module.
  findJTypeNames :: CoreM JTypeNames
  findJTypeNames = do
      nameClass <- GhcPlugins.Extras.findTHName 'Class
      nameIface <- GhcPlugins.Extras.findTHName 'Iface
      nameArray <- GhcPlugins.Extras.findTHName 'Array
      nameGeneric <- GhcPlugins.Extras.findTHName 'Generic
      namePrim <- GhcPlugins.Extras.findTHName 'Prim
      nameVoid <- GhcPlugins.Extras.findTHName 'Void
      return $ JTypeNames {..}
  
  -- | Produces a java type from a Core 'Type' if the type is sufficiently
  -- instantiated and it is of kind 'JType'.
  toJavaType :: JTypeNames -> Type -> Maybe BS.ByteString
  toJavaType JTypeNames {..} t0 = BS.concat <$> go t0
    where
      go :: Type -> Maybe [BS.ByteString]
      go (TyConApp c [LitTy (StrTyLit fs)])
        | Just n <- nameClass, tyConName c == n =
          Just [substDollar $ bytesFS fs]
        | Just n <- nameIface, tyConName c == n =
          Just [substDollar $ bytesFS fs]
      go (TyConApp c [t])
        | Just n <- nameArray, tyConName c == n =
          (++ ["[]"]) <$> go t
      go (TyConApp c [t, ts])
        | Just n <- nameGeneric, tyConName c == n = do
          bs <- go t
          args_ts <- listGo ts
          Just $ bs ++ "<" : concat (intersperse [","] args_ts) ++ [">"]
      go (TyConApp c [])
        | Just n <- nameVoid, tyConName c == n =
          Just ["void"]
      go (TyConApp c [LitTy (StrTyLit fs)])
        | Just n <- namePrim, tyConName c == n =
          Just [bytesFS fs]
      go _ = Nothing
  
      listGo :: Type -> Maybe [[BS.ByteString]]
      listGo (TyConApp c [_]) | nilDataConName == tyConName c = Just []
      listGo (TyConApp c [_, tx, txs]) | consDataConName == tyConName c =
        (:) <$> go tx <*> listGo txs
      listGo _ = Nothing
  
      -- Substitutes '$' with '.' in java names.
      substDollar :: BS.ByteString -> BS.ByteString
      substDollar xs
        | fromIntegral (ord '$') `BS.elem` xs =
          let subst (chr . fromIntegral -> '$') = fromIntegral (ord '.')
              subst x = x
           in BS.map subst xs
        | otherwise = xs
  
  -- | An occurrence of a java quasiquotation
  data QQOcc = QQOcc
      { -- | The type of the result
        qqOccResTy :: Type
        -- | The type of the arguments
      , qqOccArgTys :: [Type]
        -- | The input of the quasiquoter
      , qqOccInput :: BS.ByteString
        -- | The name of the method to generate
      , qqOccMName :: BS.ByteString
        -- | The antiquotations of the quasiquoter
      , qqOccAntiQs :: BS.ByteString
        -- | The line number where the quasiquotation appears
      , qqOccLineNumber :: Integer
      }
  
  -- | A monad for collecting qqMarker occurrences.
  type QQJavaM a = WriterT (Endo [QQOcc]) CoreM a
  
  -- Collects the occurrences of qqMarkers.
  --
  -- The program is expected to have 'qqMarker' occurrences inserted
  -- by the java quasiquoter.
  --
  -- > module A where
  -- > import Language.Java.Inline
  -- >
  -- > f = ...
  -- >     (qqMarker ... (callStatic "Inline__<pkg>_<mod>" "inline__method_i" []))
  -- >     ...
  -- >
  -- > g = ...
  -- > ...
  --
  -- 'collectQQMarkers' yields one 'QQOcc' value for every occurrence of
  -- 'qqMarker', and the program resulting from removing the markers.
  --
  -- > module A where
  -- > import Language.Java.Inline
  -- >
  -- > f = ...
  -- >     (callStatic "Inline__<pkg>_<mod>" "inline__method_i" [])
  -- >     ...
  -- >
  -- > g = ...
  -- > ...
  --
  collectQQMarkers
    :: [Name] -> CoreProgram -> CoreM (CoreProgram, [QQOcc])
  collectQQMarkers qqMarkerNames p0 = do
      (p1, e) <- runWriterT (mapM bindMarkers p0)
      return (p1, appEndo e [])
    where
      bindMarkers :: CoreBind -> QQJavaM CoreBind
      bindMarkers (NonRec b e) = NonRec b <$> expMarkers e
      bindMarkers (Rec bs) = Rec <$> mapM (\(b, e) -> (,) b <$> expMarkers e) bs
  
      expMarkers :: CoreExpr -> QQJavaM CoreExpr
      expMarkers (App (App (App (App (App (App (App (App (App (App (App (App (App
                   (App (App (App (App (App (App (App (App (Var fid) _)
                   (Type (parseArgTys -> Just tyargs)))
                   (Type tyres))
                   (Type (LitTy (StrTyLit fs_input))))
                   (Type (LitTy (StrTyLit fs_mname))))
                   (Type (LitTy (StrTyLit fs_antiqs))))
                   (Type (LitTy (NumTyLit lineNumber))))
                   _) _) _) _) _) _) _) _) _) _) _) args) _)
                   e
                 )
          | elem (idName fid) qqMarkerNames = do
        tell $ Endo $ (:) $ QQOcc
          { qqOccResTy = tyres
          , qqOccArgTys = tyargs
          , qqOccInput = bytesFS fs_input
          , qqOccMName = bytesFS fs_mname
          , qqOccAntiQs = bytesFS fs_antiqs
          , qqOccLineNumber = lineNumber
          }
        return (App e args)
      expMarkers (Var fid) | elem (idName fid) qqMarkerNames =
        lift $ GhcPlugins.Extras.failWith $
        text "inline-java Plugin: found invalid qqMarker."
      expMarkers (App e a) = App <$> expMarkers e <*> expMarkers a
      expMarkers (Lam b e) = Lam b <$> expMarkers e
      expMarkers (Let bnd e) = Let <$> bindMarkers bnd <*> expMarkers e
      expMarkers (Case e0 b t alts) = do
        e0' <- expMarkers e0
        let expAlt (Alt a bs e) = Alt a bs <$> expMarkers e
        alts' <- mapM expAlt alts
        return (Case e0' b t alts')
      expMarkers (Cast e c) = flip Cast c <$> expMarkers e
      expMarkers (Tick t e) = Tick t <$> expMarkers e
      expMarkers e@(Coercion {}) = return e
      expMarkers e@(Lit {}) = return e
      expMarkers e@(Var {}) = return e
      expMarkers e@(Type {}) = return e
  
      -- @parseArgTys (t_1, (t_2, ... (t_n, ()) ... )) = Just [t_1, t_2, ... t_n]@
      --
      -- Yields @Nothing@ when the input is not of the expected form.
      parseArgTys :: Type -> Maybe [Type]
      parseArgTys (TyConApp c [_, _, ty, tys])
        | c == promotedTupleDataCon Boxed 2 =
          (ty:) <$> parseArgTys tys
      parseArgTys (TyConApp c [])
        | c == unitTyCon =
          Just []
      parseArgTys _ = Nothing
  
  -- | Produces static structures which contain the class names and
  -- the bytecodes. For instance:
  --
  -- > static unsigned char bc0[] = {202, 254, ... }
  -- >
  -- > static unsigned char bc1[] = {202, 254, ... }
  -- >
  -- > static struct inline_java_dot_class dcs[] =
  -- >   { { "io.tweag.inlinejava.Inline__main_Language_Java_InlineSpec"
  -- >     , 2941 // length of bc0
  -- >     , bc0
  -- >     }
  -- >  , { "io.tweag.inlinejava.Inline__main_Language_Java_InlineSpec$1Foo"
  -- >    , 579 // length of bc1
  -- >    , bc1
  -- >    }
  -- > };
  -- >
  dotClasses :: [DotClass] -> SDoc
  dotClasses dcs = vcat $
        text "static int dc_count =" <+> int (length dcs) <> semi
      : [ vcat
          [ text "static unsigned char bc" <> int i <> text "[] ="
          , braces (pprWithCommas (text . show) (BS.unpack bc)) <> semi
          ]
        | (i, DotClass _ bc) <- zip [0..] dcs
        ]
        ++
        [ text "static struct inline_java_dot_class dcs[] ="
        , braces
            (pprWithCommas
              (\(i, DotClass name bc) ->
                 braces $ pprWithCommas id
                   [ text (show name), int (BS.length bc), text "bc" <> int i])
              (zip [0..] dcs)
            ) <> semi
        ]
  
  -- | Produces the constructor function which inserts the static structures
  -- generated by 'dotClasses' into the bytecode table.
  cConstructor :: DynFlags -> Module -> SDoc -> CStub
  cConstructor dflags this_mod decls =
      initializerCStub
        (targetPlatform dflags)
        (mkInitializerStubLabel this_mod "inline_java_bctable")
        decls
        (text "{ inline_java_bctable = inline_java_new_pack(inline_java_bctable, dcs, dc_count); }")
  