from-string-1.0.0.2: codegen/Main.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

module Main (main) where

import Control.Exception
import Control.Monad (unless)
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import Data.ByteString.Builder (Builder)
import qualified Data.ByteString.Builder as B
import Data.List (intercalate, intersperse)
import Data.Maybe (fromMaybe)
import RawFilePath
import System.Exit
import System.IO
import System.Posix.Env.ByteString
import System.Posix.Temp.ByteString

data StringType
  = String
  | StrictByteString
  | LazyByteString
  | ByteStringBuilder
  | ShortByteString
  | StrictText
  | LazyText
  | TextBuilder
  deriving (Eq, Show)

stringTypeName :: StringType -> Builder
stringTypeName = \case
  String -> "String"
  StrictByteString -> "B.ByteString"
  LazyByteString -> "LB.ByteString"
  ByteStringBuilder -> "BB.Builder"
  ShortByteString -> "SB.ShortByteString"
  StrictText -> "T.Text"
  LazyText -> "LT.Text"
  TextBuilder -> "TB.Builder"

data Strategy
  = ConvFunc ByteString
  | Via StringType

type Row = (StringType, StringType, Strategy)

mappings :: [Row]
mappings =
  [ (String, StrictByteString, Via StrictText)
  , (String, LazyByteString, Via LazyText)
  , (String, ByteStringBuilder, ConvFunc "BB.stringUtf8")
  , (String, ShortByteString, Via StrictByteString)
  , (String, StrictText, ConvFunc "T.pack")
  , (String, LazyText, ConvFunc "LT.pack")
  , (String, TextBuilder, ConvFunc "TB.fromString")
  , (StrictByteString, String, Via StrictText)
  , (StrictByteString, LazyByteString, ConvFunc "LB.fromStrict")
  , (StrictByteString, ByteStringBuilder, ConvFunc "BB.byteString")
  , (StrictByteString, ShortByteString, ConvFunc "SB.toShort")
  , (StrictByteString, StrictText, ConvFunc "TE.decodeUtf8Lenient")
  , (StrictByteString, LazyText, Via LazyByteString)
  , (StrictByteString, TextBuilder, Via StrictText)
  , (LazyByteString, String, Via LazyText)
  , (LazyByteString, StrictByteString, ConvFunc "LB.toStrict")
  , (LazyByteString, ByteStringBuilder, ConvFunc "BB.lazyByteString")
  , (LazyByteString, ShortByteString, Via StrictByteString)
  , (LazyByteString, StrictText, Via StrictByteString)
  , (LazyByteString, LazyText, ConvFunc "LTE.decodeUtf8With TEE.lenientDecode")
  , (LazyByteString, TextBuilder, Via LazyText)
  , (ByteStringBuilder, String, Via LazyText)
  , (ByteStringBuilder, StrictByteString, Via LazyByteString)
  , (ByteStringBuilder, LazyByteString, ConvFunc "BB.toLazyByteString")
  , (ByteStringBuilder, ShortByteString, Via StrictByteString)
  , (ByteStringBuilder, StrictText, Via StrictByteString)
  , (ByteStringBuilder, LazyText, Via LazyByteString)
  , (ByteStringBuilder, TextBuilder, Via LazyByteString)
  , (ShortByteString, String, Via StrictByteString)
  , (ShortByteString, StrictByteString, ConvFunc "SB.fromShort")
  , (ShortByteString, LazyByteString, Via StrictByteString)
  , (ShortByteString, ByteStringBuilder, ConvFunc "BB.shortByteString")
  , (ShortByteString, StrictText, Via StrictByteString)
  , (ShortByteString, LazyText, Via StrictByteString)
  , (ShortByteString, TextBuilder, Via StrictByteString)
  , (StrictText, String, ConvFunc "T.unpack")
  , (StrictText, StrictByteString, ConvFunc "TE.encodeUtf8")
  , (StrictText, LazyByteString, Via LazyText)
  , (StrictText, ByteStringBuilder, ConvFunc "TE.encodeUtf8Builder")
  , (StrictText, ShortByteString, Via StrictByteString)
  , (StrictText, LazyText, ConvFunc "LT.fromStrict")
  , (StrictText, TextBuilder, ConvFunc "TB.fromText")
  , (LazyText, String, ConvFunc "LT.unpack")
  , (LazyText, StrictByteString, Via LazyByteString)
  , (LazyText, LazyByteString, ConvFunc "LTE.encodeUtf8")
  , (LazyText, ByteStringBuilder, ConvFunc "LTE.encodeUtf8Builder")
  , (LazyText, ShortByteString, Via StrictByteString)
  , (LazyText, StrictText, ConvFunc "LT.toStrict")
  , (LazyText, TextBuilder, ConvFunc "TB.fromLazyText")
  , (TextBuilder, String, Via LazyText)
  , (TextBuilder, StrictByteString, Via LazyText)
  , (TextBuilder, LazyByteString, Via LazyText)
  , (TextBuilder, ByteStringBuilder, Via LazyText)
  , (TextBuilder, ShortByteString, Via LazyText)
  , (TextBuilder, StrictText, Via LazyText)
  , (TextBuilder, LazyText, ConvFunc "TB.toLazyText")
  ]

fullyQualifiedModule :: ByteString -> ByteString
fullyQualifiedModule = \case
  "B" -> "Data.ByteString"
  "BB" -> "Data.ByteString.Builder"
  "LB" -> "Data.ByteString.Lazy"
  "SB" -> "Data.ByteString.Short"
  "T" -> "Data.Text"
  "TE" -> "Data.Text.Encoding"
  "TEE" -> "Data.Text.Encoding.Error"
  "LT" -> "Data.Text.Lazy"
  "TB" -> "Data.Text.Lazy.Builder"
  "LTE" -> "Data.Text.Lazy.Encoding"
  unknown -> error $ show unknown <> " is not a correct module name"

fullyQualifiedExpr :: ByteString -> ByteString
fullyQualifiedExpr expr = fullyQualifiedModule moduleName <> rest
 where
  (moduleName, rest) = B.break (0x2e ==) expr

forceLookup :: (Show k, Eq k) => k -> [(k, v)] -> v
forceLookup key table = fromMaybe (error $ "key " <> show key <> " not found") $ lookup key table

getStrategy
  :: StringType -> StringType -> [Row] -> Strategy
getStrategy fromType toType table = forceLookup (fromType, toType) kvTable
 where
  kvTable = flip map table $ \(a, b, c) -> ((a, b), c)

getMapping
  :: Row -> [Row] -> [ByteString]
getMapping (fromType, toType, strategy) table = case strategy of
  ConvFunc expr -> [expr]
  Via intermediate ->
    let formerStrategy = getStrategy fromType intermediate table
        formerPart = getMapping (fromType, intermediate, formerStrategy) table
        latterStrategy = getStrategy intermediate toType table
        latterPart = getMapping (intermediate, toType, latterStrategy) table
     in latterPart ++ formerPart

main :: IO ()
main = do
  args <- getArgs
  let noDiff = "--no-diff" `elem` args
      shouldKeepFile = noDiff || "--keep" `elem` args
  bracket acquire (release shouldKeepFile) $ \(tmpPath, hdl) -> do
    B.putStr $ tmpPath <> "\n"
    B.hPutBuilder hdl $ header <> mconcat instancePart
    hClose hdl

    unless noDiff $ do
      diffExitCode <- runDiff tmpPath
      exitWith diffExitCode
 where
  acquire = mkstemps "haskell-from-" ".hs"
  release shouldKeepFile (path, hdl)
    | shouldKeepFile = hClose hdl
    | otherwise = hClose hdl >> removeFile path

  runDiff path =
    runProcess
      "diff"
      ["-u", "--color=always", path, "src/From/String/AutoGen.hs"]

  instanceLine fromType toType =
    mconcat
      [ "instance From "
      , stringTypeName fromType
      , " "
      , stringTypeName toType
      , " where\n"
      ]
  instancePart =
    [ let exprs = getMapping row mappings
          quote b = "'" <> b <> "'"
       in mconcat
            [ "\n"
            , "-- | "
            , mconcat
                $ intersperse " "
                $ map
                  (quote . B.byteString)
                $ intercalate ["."]
                -- Split by space, to link each part in the expression
                $ map (map fullyQualifiedExpr . B.split 0x20) exprs
            , "\n"
            , instanceLine fromType toType
            , "  from = "
            , mconcat
                $ intersperse
                  " . "
                $ map
                  B.byteString
                  exprs
            , "\n"
            ]
    | row@(fromType, toType, _) <- mappings
    ]

runProcess :: RawFilePath -> [ByteString] -> IO ExitCode
runProcess cmd args = startProcess (proc cmd args) >>= waitForProcess

header :: Builder
header =
  "{- FOURMOLU_DISABLE -}\n\
  \{- Make sure you edit test/Main.hs -}\n\
  \{-# LANGUAGE FlexibleInstances #-}\n\
  \{-# LANGUAGE MultiParamTypeClasses #-}\n\
  \{-# OPTIONS_GHC -Wno-orphans #-}\n\
  \\n\
  \module From.String.AutoGen () where\n\
  \\n\
  \import qualified Data.ByteString as B\n\
  \import qualified Data.ByteString.Builder as BB\n\
  \import qualified Data.ByteString.Lazy as LB\n\
  \import qualified Data.ByteString.Short as SB\n\
  \import qualified Data.Text as T\n\
  \import qualified Data.Text.Encoding as TE\n\
  \import qualified Data.Text.Encoding.Error as TEE\n\
  \import qualified Data.Text.Lazy as LT\n\
  \import qualified Data.Text.Lazy.Builder as TB\n\
  \import qualified Data.Text.Lazy.Encoding as LTE\n\
  \import From (From (..))\n"