THCompareSpec.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{- HLINT ignore "Use fewer imports" -}

module Test.MockCat.THCompareSpec (spec) where

import Test.Hspec
import Language.Haskell.TH (pprint, Dec(..), Type(..), litE, stringL, listE, tupE)
import Language.Haskell.TH.Syntax (nameBase)
import Data.Char (isDigit, isLower)
import Data.List (sort, nub)
import qualified Data.List as L
import Data.Maybe (mapMaybe)
import qualified Data.Map.Strict as Map
import qualified Data.Text as T
import qualified System.IO as SIO
import Test.MockCat.SharedSpecDefs
import Test.MockCat.TH (makeMock, makeAutoLiftMock, makePartialMock, makeAutoLiftPartialMock)
import Language.Haskell.TH (lookupTypeName, conT)
import Control.Monad (forM_, when)

normalize :: String -> String
normalize =
  T.unpack
    . squeezePunctuation
    . T.unwords
    . map stripVarSuffix
    . T.words
    . stripQualifiers
    . T.pack

stripQualifiers :: T.Text -> T.Text
stripQualifiers =
  stripPrefixes
    (map
      T.pack
      [ "Test.MockCat.SharedSpecDefs."
      , "Test.MockCat.MockT."
      , "Test.MockCat.Internal.Builder."
      , "Test.MockCat.Cons."
      , "Test.MockCat.Param."
      , "Control.Monad.IO.Class."
      , "GHC.Types."
      , "Data.Typeable.Internal."
      , "GHC.Internal."
      , "Base."
      , "Test.MockCat.Verify."
      , "Verify."
      , "GHC.Show."
      , "GHC.Classes."
      , "Show."
      , "Classes."
      , "Internal."
      , "Test.MockCat.Mock."
      , "Test.MockCat.Types."
      , "Test.MockCat.WithMock."
      ])
  where
    stripPrefixes [] t = t
    stripPrefixes (p : ps) t = stripPrefixes ps (T.replace p T.empty t)

stripVarSuffix :: T.Text -> T.Text
stripVarSuffix = T.pack . go Nothing . T.unpack
  where
    go _ [] = []
    go prev (c : cs)
      | c == '_' && maybe False isLower prev =
          go prev (dropWhile isDigit cs)
      | otherwise = c : go (Just c) cs

squeezePunctuation :: T.Text -> T.Text
squeezePunctuation =
  applyReplacements
    [ ("( ", "(")
    , (" )", ")")
    , (" ,", ",")
    , ("[ ", "[")
    , (" ]", "]")
    , (" <" , "<")
    , ("> ", ">")
    , ("((:>) ", "(")
    , ("((:>)", "(")
    , ("(:>) ", "")
    , ("(:>)", "")
    , ("([a])", "[a]")
    ]
  . T.replace (T.pack "Head (Param") (T.pack "Head :>Param")
  . T.replace (T.pack "Head :> Param") (T.pack "Head :>Param")
  -- Fix: MockBuilder (Head :>Param r) r () should not become MockBuilder (Head :>Param r r ()
  -- First fix the case where " r r ()" appears (missing closing paren)
  . T.replace (T.pack ":>Param r r ()") (T.pack ":>Param r) r ()")
  . T.replace (T.pack "Head :>Param r r ()") (T.pack "Head :>Param r) r ()")
  . T.replace (T.pack "Param r r ()") (T.pack "Param r) r ()")
  -- Then handle the general case of ") r ()" -> " r ()" but only if not already fixed
  . T.replace (T.pack ") r ()") (T.pack " r ()")
  . T.replace (T.pack "Typeable m,") T.empty
  . T.replace (T.pack ", Typeable m") T.empty
  . T.replace (T.pack "Typeable a,") T.empty
  . T.replace (T.pack ", Typeable a") T.empty
  . T.replace (T.pack " )") (T.pack ")")
  . T.replace (T.pack "Typeable [a]") (T.pack "Typeable a")
  . T.replace (T.pack "Typeable ([a])") (T.pack "Typeable a")
  . T.replace (T.pack "Typeable (ResultType m)") (T.pack "Typeable m")
  where
    applyReplacements reps txt =
      L.foldl'
        (\acc (from, to) -> T.replace (T.pack from) (T.pack to) acc)
        txt
        reps

-- generate at compile time to avoid running restricted TH actions in IO
-- store pretty-printed generated declarations as strings
generatedFinderStr :: String
generatedFinderStr = $(do decs <- makeAutoLiftPartialMock [t|Finder|]; litE (stringL (concatMap pprint decs)))

generatedFinderSigPairs :: [(String, String)]
generatedFinderSigPairs =
  $(
    do
      decs <- makeAutoLiftPartialMock [t|Finder|]
      let sigPairs =
            [ (nameBase name, pprint ty)
            | SigD name ty <- decs
            ]
      listE
        [ tupE [litE (stringL fn), litE (stringL sig)]
        | (fn, sig) <- sigPairs
        ]
  )

generatedFinderSigMap :: Map.Map String String
generatedFinderSigMap = Map.fromList generatedFinderSigPairs

generatedUserInputSigPairs :: [(String, String)]
generatedUserInputSigPairs =
  $(
    do
      decs <- makeAutoLiftPartialMock [t|UserInputGetter|]
      let sigPairs =
            [ (nameBase name, pprint ty)
            | SigD name ty <- decs
            ]
      listE
        [ tupE [litE (stringL fn), litE (stringL sig)]
        | (fn, sig) <- sigPairs
        ]
  )

generatedUserInputSigMap :: Map.Map String String
generatedUserInputSigMap = Map.fromList generatedUserInputSigPairs

generatedFileOpStr :: String
generatedFileOpStr = $(do decs <- makeAutoLiftPartialMock [t|FileOperation|]; litE (stringL (concatMap pprint decs)))

generatedTestClassStr :: String
generatedTestClassStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.TestClass"
    case m of
      Nothing -> fail "TestClass not found"
      Just n -> do decs <- makeAutoLiftMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedMultiApplyStr :: String
generatedMultiApplyStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.MultiApplyTest"
    case m of
      Nothing -> fail "MultiApplyTest not found"
      Just n -> do decs <- makeAutoLiftMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedExplicitStr :: String
generatedExplicitStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.ExplicitlyReturnMonadicValuesTest"
    case m of
      Nothing -> fail "ExplicitlyReturnMonadicValuesTest not found"
      Just n -> do decs <- makeMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedDefaultMethodStr :: String
generatedDefaultMethodStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.DefaultMethodTest"
    case m of
      Nothing -> fail "DefaultMethodTest not found"
      Just n -> do decs <- makeAutoLiftMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedAssocTypeStr :: String
generatedAssocTypeStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.AssocTypeTest"
    case m of
      Nothing -> fail "AssocTypeTest not found"
      Just n -> do decs <- makeAutoLiftMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedAssocSigPairs :: [(String, String)]
generatedAssocSigPairs =
  $(
    do
      m <- lookupTypeName "Test.MockCat.SharedSpecDefs.AssocTypeTest"
      case m of
        Nothing -> fail "AssocTypeTest not found"
        Just n -> do
          decs <- makeAutoLiftMock (conT n)
          let sigPairs =
                [ (nameBase name, pprint ty)
                | SigD name ty <- decs
                ]
          listE
            [ tupE [litE (stringL fn), litE (stringL sig)]
            | (fn, sig) <- sigPairs
            ]
  )

generatedAssocSigMap :: Map.Map String String
generatedAssocSigMap = Map.fromList generatedAssocSigPairs

generatedParamThreeStr :: String
generatedParamThreeStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.ParamThreeMonad"
    case m of
      Nothing -> fail "ParamThreeMonad not found"
      Just n -> do
        let targetType =
              AppT
                (AppT (ConT n) (ConT ''Int))
                (ConT ''Bool)
        decs <- makeAutoLiftMock (pure targetType)
        litE (stringL (concatMap pprint decs))
  )

generatedUserInputStr :: String
generatedUserInputStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.UserInputGetter"
    case m of
      Nothing -> fail "UserInputGetter not found"
      Just n -> do decs <- makeAutoLiftPartialMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedExplicitPartialStr :: String
generatedExplicitPartialStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.ExplicitlyReturnMonadicValuesPartialTest"
    case m of
      Nothing -> fail "ExplicitlyReturnMonadicValuesPartialTest not found"
      Just n -> do decs <- makePartialMock (conT n); litE (stringL (concatMap pprint decs))
  )

-- additional generated declarations for classes produced in `TypeClassTHSpec.hs`


generatedVar2_1SubStr :: String
generatedVar2_1SubStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.MonadVar2_1Sub"
    case m of
      Nothing -> fail "MonadVar2_1Sub not found"
      Just n -> do decs <- makeAutoLiftMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedVar2_2SubStr :: String
generatedVar2_2SubStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.MonadVar2_2Sub"
    case m of
      Nothing -> fail "MonadVar2_2Sub not found"
      Just n -> do decs <- makeAutoLiftMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedVar3_1SubStr :: String
generatedVar3_1SubStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.MonadVar3_1Sub"
    case m of
      Nothing -> fail "MonadVar3_1Sub not found"
      Just n -> do decs <- makeAutoLiftMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedVar3_2SubStr :: String
generatedVar3_2SubStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.MonadVar3_2Sub"
    case m of
      Nothing -> fail "MonadVar3_2Sub not found"
      Just n -> do decs <- makeAutoLiftMock (conT n); litE (stringL (concatMap pprint decs))
  )

generatedVar3_3SubStr :: String
generatedVar3_3SubStr = $(
  do
    m <- lookupTypeName "Test.MockCat.SharedSpecDefs.MonadVar3_3Sub"
    case m of
      Nothing -> fail "MonadVar3_3Sub not found"
      Just n -> do decs <- makeAutoLiftMock (conT n); litE (stringL (concatMap pprint decs))
  )

partialMockSpecPath :: FilePath
partialMockSpecPath = "test/Test/MockCat/PartialMockSpec.hs"

typeClassSpecPath :: FilePath
typeClassSpecPath = "test/Test/MockCat/TypeClassSpec.hs"

-- extract textual instance context for handwritten module
extractHandwrittenInstanceCxts :: FilePath -> String -> IO [String]
extractHandwrittenInstanceCxts fp className = do
  src <- SIO.readFile fp
  pure $ extractInstanceHeaders src className

spec :: Spec
spec = describe "TH generated vs handwritten instances" do
  let cases =
        [ ("Finder partial mock instance matches handwritten", partialMockSpecPath, "Finder", generatedFinderStr)
        , ("FileOperation mock instance matches handwritten", partialMockSpecPath, "FileOperation", generatedFileOpStr)
        , ("UserInputGetter partial mock instance matches handwritten", partialMockSpecPath, "UserInputGetter", generatedUserInputStr)
        , ("ExplicitlyReturnMonadicValuesPartialTest instance matches handwritten", partialMockSpecPath, "ExplicitlyReturnMonadicValuesPartialTest", generatedExplicitPartialStr)
        , ("TypeClass: TestClass instance matches handwritten", typeClassSpecPath, "TestClass", generatedTestClassStr)
        , ("TypeClass: MultiApplyTest instance matches handwritten", typeClassSpecPath, "MultiApplyTest", generatedMultiApplyStr)
        , ("TypeClass: ExplicitlyReturnMonadicValuesTest instance matches handwritten", typeClassSpecPath, "ExplicitlyReturnMonadicValuesTest", generatedExplicitStr)
        , ("TypeClass: DefaultMethodTest instance matches handwritten", typeClassSpecPath, "DefaultMethodTest", generatedDefaultMethodStr)
        , ("TypeClass: AssocTypeTest instance matches handwritten", typeClassSpecPath, "AssocTypeTest", generatedAssocTypeStr)
        , ("TypeClass: ParamThreeMonad instance matches handwritten", typeClassSpecPath, "ParamThreeMonad", generatedParamThreeStr)

        , ("MonadVar2_1Sub instance matches handwritten", typeClassSpecPath, "MonadVar2_1Sub", generatedVar2_1SubStr)
        , ("MonadVar2_2Sub instance matches handwritten", typeClassSpecPath, "MonadVar2_2Sub", generatedVar2_2SubStr)
        , ("MonadVar3_1Sub instance matches handwritten", typeClassSpecPath, "MonadVar3_1Sub", generatedVar3_1SubStr)
        , ("MonadVar3_2Sub instance matches handwritten", typeClassSpecPath, "MonadVar3_2Sub", generatedVar3_2SubStr)
        , ("MonadVar3_3Sub instance matches handwritten", typeClassSpecPath, "MonadVar3_3Sub", generatedVar3_3SubStr)
        ]
  forM_ cases \(desc, manualPath, className, generatedStr) ->
    it desc $
      assertInstancesEquivalent manualPath className generatedStr

  it "generated helper signatures do not contain duplicated Typeable constraints" do
    let allSigMaps =
          [ generatedFinderSigMap,
            generatedUserInputSigMap
          ]
        extractTypeableTargets :: String -> [String]
        extractTypeableTargets s =
          let go acc txt =
                let (_, rest) = T.breakOn (T.pack "Typeable") txt
                 in if T.null rest
                      then acc
                      else
                        let afterTypeable = T.drop (T.length (T.pack "Typeable")) rest
                            afterTrim = T.stripStart afterTypeable
                            tok = T.strip $ T.takeWhile (\c -> c /= ',' && c /= ')' && c /= '=') afterTrim
                            rest' = afterTypeable
                         in go (acc ++ [T.unpack tok]) rest'
           in go [] (T.pack s)
        checkMap m =
          forM_ (Map.toList m) \(_fn, sig) -> do
            let nsig = normalizeSignature sig
                targets = extractTypeableTargets nsig
            length targets `shouldBe` length (nub targets)
    mapM_ checkMap allSigMaps

  describe "Partial mock helper signatures match handwritten" do
    it "_findIds signature matches handwritten" $
      assertHelperSigMatches partialMockSpecPath generatedFinderSigMap "_findIds"

    it "_toUserInput signature contains ResolvableParamsOf (due to MockResult wrapper)" do
      let genSig =
            Map.lookup "_toUserInput" generatedUserInputSigMap
      case genSig of
        Nothing ->
          expectationFailure "TH generated signature not found: _toUserInput"
        Just sig ->
          normalizeSignature sig
            `shouldSatisfy` T.isInfixOf (T.pack "ResolvableParamsOf") . T.pack

    it "_produce signature matches handwritten" $
      assertHelperSigMatches typeClassSpecPath generatedAssocSigMap "_produce"

-- helper to extract generated instance headers from a pre-pprinted string
extractGeneratedInstanceCxtsFromStr :: String -> String -> [String]
extractGeneratedInstanceCxtsFromStr = extractInstanceHeaders

extractInstanceHeaders :: String -> String -> [String]
extractInstanceHeaders input className =
  let txt = T.pack input
      classTxt = T.pack className
      insts = drop 1 $ T.splitOn (T.pack "instance") txt
      mkHeader body =
        let (beforeWhere, _) = T.breakOn (T.pack "where") body
            (beforeBrace, _) = T.breakOn (T.pack "{") beforeWhere
            header = T.strip beforeBrace
            headerWithoutQualifiers = stripQualifiers header
            headPart =
              let (_ctxPart, rest) = T.breakOn (T.pack "=>") headerWithoutQualifiers
               in if T.null rest
                    then headerWithoutQualifiers
                    else T.strip $ T.drop 2 rest
            actualClassName =
              case T.words headPart of
                [] -> Nothing
                (tok : _) ->
                  let cleaned = T.dropAround (`elem` "()") tok
                   in if T.null cleaned then Nothing else Just cleaned
            containsMockT = T.pack "MockT" `T.isInfixOf` headerWithoutQualifiers
         in case actualClassName of
              Just cls
                | cls == classTxt && containsMockT ->
                    Just $ normalize $ "instance " ++ T.unpack headerWithoutQualifiers
              _ -> Nothing
   in mapMaybe mkHeader insts

assertInstancesEquivalent :: FilePath -> String -> String -> Expectation
assertInstancesEquivalent manualPath className generatedStr = do
  hand <- extractHandwrittenInstanceCxts manualPath className
  let gen = extractGeneratedInstanceCxtsFromStr generatedStr className
      sortedHand = sort hand
      sortedGen = sort gen
  when (null hand) $ expectationFailure $
      "Handwritten instance not found: " ++ className ++ " in " ++ manualPath
  when (null gen) $ expectationFailure $
      "TH generated instance not found: " ++ className
        ++ "\nGenerated result (prefix only):\n"
        ++ take 1000 generatedStr
  sortedGen `shouldBe` sortedHand

extractHandwrittenFunctionSig :: FilePath -> String -> IO (Maybe String)
extractHandwrittenFunctionSig fp funName = do
  src <- SIO.readFile fp
  let linesText = T.lines (T.pack src)
      isSignatureLine line =
        let trimmed = T.stripStart line
         in (T.pack funName <> T.pack " ::") `T.isPrefixOf` trimmed
      continuation line =
        case T.uncons line of
          Just (c, _) -> c == ' ' || c == '\t'
          Nothing -> False
      collect [] = Nothing
      collect (line:rest)
        | isSignatureLine line =
            let cont = takeWhile continuation rest
             in Just $ T.unpack $ T.unlines (line : cont)
        | otherwise = collect rest
  pure $ collect linesText

stripForallClause :: String -> String
stripForallClause sig =
  let txt = T.pack sig
      (preForall, rest) = T.breakOn (T.pack "forall") txt
   in
    if T.null rest
      then sig
      else
        let afterForall = T.drop 1 $ T.dropWhile (/= '.') rest
         in T.unpack (preForall <> afterForall)

normalizeSignature :: String -> String
normalizeSignature =
  T.unpack
    . T.strip
    . T.pack
    . dropFunctionName
    . stripForallClause
    . replaceDoubleParens
    . normalize

-- collapse duplicated closing parens introduced by normalization quirks
replaceDoubleParens :: String -> String
replaceDoubleParens = T.unpack . T.replace (T.pack "))") (T.pack ")") . T.pack

dropFunctionName :: String -> String
dropFunctionName sig =
  let txt = T.pack sig
      (_, rest) = T.breakOn (T.pack "::") txt
   in
    if T.null rest
      then sig
      else T.unpack $ T.drop 2 rest

assertHelperSigMatches ::
  FilePath ->
  Map.Map String String ->
  String ->
  Expectation
assertHelperSigMatches manualPath generatedSigMap funName = do
  manualSigM <- extractHandwrittenFunctionSig manualPath funName
  manualSig <-
    case manualSigM of
      Nothing -> do
        expectationFailure $
          "Handwritten signature not found: "
            ++ funName
            ++ " in "
            ++ manualPath
        pure ""
      Just sig -> pure sig
  generatedSig <-
    case Map.lookup funName generatedSigMap of
      Nothing -> do
        expectationFailure $
          "TH generated signature not found: "
            ++ funName
        pure ""
      Just sig -> pure sig
  normalizeSignature generatedSig `shouldBe` normalizeSignature manualSig