Type.hs

{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Test.Properties.Arb.Type
  ( genType,
    shrinkType,
    shrinkTyVarBinders,
    shrinkForallTelescope,
  )
where

import Aihc.Parser.Syntax
import Data.Text (Text)
import Data.Text qualified as T
import Test.Properties.Arb.Expr (shrinkExpr)
import Test.Properties.Arb.Identifiers
  ( genCharValue,
    genConName,
    genQuasiBody,
    genQuoterName,
    genVarId,
    genVarIdNoHash,
    shrinkIdent,
    shrinkName,
  )
import Test.Properties.Arb.Utils (smallList0)
import Test.QuickCheck

instance Arbitrary Type where
  arbitrary = genType
  shrink = shrinkType

-- | Generate a random type. Uses QuickCheck's size parameter to control
-- recursion depth.
genType :: Gen Type
genType = scale (`div` 2) $ do
  n <- getSize
  if n <= 0
    then
      oneof
        [ TVar <$> genTypeVarName,
          (`TCon` Unpromoted) <$> genConName,
          (`TCon` Promoted) <$> genConName,
          TBuiltinCon <$> genTypeBuiltinCon,
          TTypeLit <$> genTypeLiteral,
          pure (TStar "*"),
          pure TWildcard,
          TQuasiQuote <$> genQuoterName <*> genQuasiBody,
          TTuple Boxed Unpromoted <$> elements [[], [TVar "a", TCon (qualifyName Nothing (mkUnqualifiedName NameConId "B")) Unpromoted]],
          TTuple Unboxed Unpromoted <$> elements [[], [TVar "a", TCon (qualifyName Nothing (mkUnqualifiedName NameConId "B")) Unpromoted]],
          TList Unpromoted <$> genTypeListElems,
          TUnboxedSum <$> genUnboxedSumElems
        ]
    else
      oneof
        [ TVar <$> genTypeVarName,
          (`TCon` Unpromoted) <$> genConName,
          (`TCon` Promoted) <$> genConName,
          TBuiltinCon <$> genTypeBuiltinCon,
          TTypeLit <$> genTypeLiteral,
          pure (TStar "*"),
          pure TWildcard,
          TQuasiQuote <$> genQuoterName <*> genQuasiBody,
          TForall <$> genForallTelescope <*> genType,
          genTypeApp,
          genTypeTypeApp,
          genTypeInfix,
          genTypeFun,
          TTuple Boxed Unpromoted <$> genTypeTupleElems,
          TTuple Boxed Promoted <$> genPromotedTupleElems,
          TTuple Unboxed Unpromoted <$> genTypeTupleElems,
          TUnboxedSum <$> genUnboxedSumElems,
          TList Unpromoted <$> genTypeListElems,
          TList Promoted <$> genPromotedListElems,
          TParen <$> genType,
          TSplice <$> genTypeSpliceBody,
          genTypeContext,
          genTypeImplicitParam,
          TKindSig <$> genKindSigSubject <*> genKindSigKind
        ]

genTypeApp :: Gen Type
genTypeApp = TApp <$> genType <*> genType

genTypeTypeApp :: Gen Type
genTypeTypeApp = TTypeApp <$> genType <*> genType

genTypeInfix :: Gen Type
genTypeInfix = do
  lhs <- genType
  rhs <- genType
  op <- genConName
  pure (TInfix lhs op Unpromoted rhs)

genTypeFun :: Gen Type
genTypeFun =
  oneof
    [ TFun ArrowUnrestricted <$> genType <*> genType,
      TFun ArrowLinear <$> genType <*> genType,
      (TFun . ArrowExplicit <$> genMultiplicityType) <*> genType <*> genType
    ]

-- | Generate a multiplicity type for explicit multiplicity annotations.
-- Keep it simple to avoid overly complex nesting.
genMultiplicityType :: Gen Type
genMultiplicityType =
  oneof
    [ TVar <$> genTypeVarName,
      (`TCon` Unpromoted) <$> genConName
    ]

-- | Generate the body of a TH type splice: either a bare variable or a parenthesized expression.
genTypeSpliceBody :: Gen Expr
genTypeSpliceBody =
  oneof
    [ EVar <$> genTypeVarExprName,
      EParen . EVar <$> genTypeVarExprName
    ]

genTypeContext :: Gen Type
genTypeContext = do
  n <- chooseInt (1, 3)
  constraints <- vectorOf n genConstraintType
  TContext constraints <$> genType

-- | Generate a constraint type (used in contexts).
-- Typically a type constructor applied to some arguments.
genConstraintType :: Gen Type
genConstraintType = do
  s <- getSize
  className <- (`TCon` Unpromoted) <$> genConName
  oneof
    [ -- Simple constraint: ClassName tyvar
      TApp className . TVar <$> genTypeVarName,
      -- Applied constraint: ClassName (Type)
      TApp className . TParen <$> resize s genType
    ]

genTypeImplicitParam :: Gen Type
genTypeImplicitParam = do
  name <- ("?" <>) <$> genVarIdNoHash
  TImplicitParam name <$> genType

genTypeTupleElems :: Gen [Type]
genTypeTupleElems = do
  isUnit <- arbitrary
  if isUnit
    then pure []
    else do
      n <- chooseInt (2, 4)
      scale (`div` n) $ vectorOf n genType

genTypeListElems :: Gen [Type]
genTypeListElems = do
  n <- chooseInt (1, 4)
  scale (`div` n) $ vectorOf n genType

genUnboxedSumElems :: Gen [Type]
genUnboxedSumElems = do
  n <- chooseInt (2, 4)
  scale (`div` n) $ vectorOf n genType

-- | Generate elements for a promoted tuple or list. Uses simple types only
-- to avoid nesting ambiguities with kind signatures and unboxed tuples
-- inside promoted containers.
genPromotedTupleElems :: Gen [Type]
genPromotedTupleElems = do
  isUnit <- arbitrary
  if isUnit
    then pure []
    else do
      n <- chooseInt (2, 3)
      vectorOf n genPromotedElem

genPromotedListElems :: Gen [Type]
genPromotedListElems = do
  n <- chooseInt (1, 3)
  vectorOf n genPromotedElem

-- | Generate a simple type suitable for use inside promoted tuples/lists.
-- Avoids character type literals since 'c' conflicts with the promotion tick.
genPromotedElem :: Gen Type
genPromotedElem =
  oneof
    [ TVar <$> genTypeVarName,
      (`TCon` Unpromoted) <$> genConName,
      genPromotedSafeTypeLiteral,
      pure (TStar "*")
    ]

-- | Generate type literals safe for use in promoted contexts (no char literals).
genPromotedSafeTypeLiteral :: Gen Type
genPromotedSafeTypeLiteral =
  TTypeLit
    <$> oneof
      [ do
          n <- chooseInteger (0, 1000)
          pure (TypeLitInteger n (T.pack (show n))),
        do
          txt <- genSymbolText
          pure (TypeLitSymbol txt (T.pack (show (T.unpack txt))))
      ]

genSimpleTypeAtom :: Gen Type
genSimpleTypeAtom =
  oneof
    [ TVar <$> genTypeVarName,
      (`TCon` Unpromoted) <$> genConName,
      TTypeLit <$> genTypeLiteral,
      pure (TStar "*"),
      pure TWildcard,
      TQuasiQuote <$> genQuoterName <*> genQuasiBody,
      TTuple Boxed Unpromoted <$> genTypeTupleElems,
      TTuple Unboxed Unpromoted <$> smallList0 genType,
      TUnboxedSum <$> genUnboxedSumElems,
      TList Unpromoted <$> genTypeListElems,
      TParen <$> genType
    ]

genKindSigSubject :: Gen Type
genKindSigSubject = genSimpleTypeAtom

genKindSigKind :: Gen Type
genKindSigKind =
  frequency
    [ (3, genSimpleTypeAtom),
      (1, TFun ArrowUnrestricted <$> genSimpleTypeAtom <*> genSimpleTypeAtom)
    ]

genForallTelescope :: Gen ForallTelescope
genForallTelescope = do
  vis <- elements [ForallInvisible, ForallVisible]
  -- Visible foralls (forall a -> T) require specified binders; inferred ({a}) are only valid in invisible foralls
  binders <- genTypeBindersFor vis
  pure (ForallTelescope vis binders)

genTypeBindersFor :: ForallVis -> Gen [TyVarBinder]
genTypeBindersFor vis = do
  n <- chooseInt (1, 3)
  vectorOf n (genTyVarBinderFor vis)

genTyVarBinderFor :: ForallVis -> Gen TyVarBinder
genTyVarBinderFor vis = do
  name <- genTypeVarName
  case vis of
    ForallVisible ->
      oneof
        [ -- Plain specified binder: a
          pure (TyVarBinder [] (renderUnqualifiedName name) Nothing TyVarBSpecified TyVarBVisible),
          -- Kinded specified binder: (a :: Kind)
          do
            kind <- resize 0 genSimpleTypeAtom
            pure (TyVarBinder [] (renderUnqualifiedName name) (Just kind) TyVarBSpecified TyVarBVisible)
        ]
    ForallInvisible ->
      oneof
        [ -- Plain specified binder: a
          pure (TyVarBinder [] (renderUnqualifiedName name) Nothing TyVarBSpecified TyVarBVisible),
          -- Plain inferred binder: {a}
          pure (TyVarBinder [] (renderUnqualifiedName name) Nothing TyVarBInferred TyVarBVisible),
          -- Kinded inferred binder: {a :: Kind}
          do
            kind <- resize 0 genSimpleTypeAtom
            pure (TyVarBinder [] (renderUnqualifiedName name) (Just kind) TyVarBInferred TyVarBVisible),
          -- Kinded specified binder: (a :: Kind)
          do
            kind <- resize 0 genSimpleTypeAtom
            pure (TyVarBinder [] (renderUnqualifiedName name) (Just kind) TyVarBSpecified TyVarBVisible)
        ]

genTypeVarName :: Gen UnqualifiedName
genTypeVarName = mkUnqualifiedName NameVarId <$> genVarId

genTypeVarExprName :: Gen Name
genTypeVarExprName = qualifyName Nothing . mkUnqualifiedName NameVarId <$> genVarId

genTypeLiteral :: Gen TypeLiteral
genTypeLiteral =
  oneof
    [ do
        n <- chooseInteger (0, 1000)
        pure (TypeLitInteger n (T.pack (show n))),
      do
        txt <- genSymbolText
        pure (TypeLitSymbol txt (T.pack (show (T.unpack txt)))),
      do
        c <- genCharValue
        pure (TypeLitChar c (T.pack (show c)))
    ]

genTypeBuiltinCon :: Gen TypeBuiltinCon
genTypeBuiltinCon =
  elements
    [ TBuiltinTuple 2,
      TBuiltinTuple 3,
      TBuiltinArrow,
      TBuiltinList,
      TBuiltinCons
    ]

genSymbolText :: Gen Text
genSymbolText = do
  len <- chooseInt (0, 8)
  chars <- vectorOf len genCharValue
  pure (T.pack chars)

shrinkType :: Type -> [Type]
shrinkType ty =
  [TWildcard | ty /= TWildcard]
    ++ case ty of
      TVar name ->
        [TVar $ unqualifiedNameFromText "a" | renderUnqualifiedName name /= "a"]
          <> [TVar (mkUnqualifiedName NameVarId shrunk) | shrunk <- shrinkIdent (renderUnqualifiedName name)]
      TCon name promoted ->
        [ TCon shrunk promoted
        | shrunk <- shrinkName name
        ]
      TBuiltinCon {} ->
        []
      TImplicitParam name inner ->
        [inner]
          <> [TImplicitParam name' inner | name' <- shrinkImplicitParamName name]
          <> [TImplicitParam name inner' | inner' <- shrinkType inner]
      TTypeLit {} ->
        []
      TStar {} ->
        []
      TQuasiQuote quoter body ->
        [TQuasiQuote q body | q <- shrinkIdent quoter]
          <> [TQuasiQuote quoter b | b <- map T.pack (shrink (T.unpack body))]
      TForall telescope inner ->
        [inner]
          <> [TForall telescope' inner | telescope' <- shrinkForallTelescope telescope]
          <> [TForall telescope inner' | inner' <- shrinkType inner]
      TApp fn arg ->
        [fn, arg]
          <> [TApp fn' arg | fn' <- shrinkType fn]
          <> [TApp fn arg' | arg' <- shrinkType arg]
      TTypeApp fn arg ->
        [fn, arg]
          <> [TTypeApp fn' arg | fn' <- shrinkType fn]
          <> [TTypeApp fn arg' | arg' <- shrinkType arg]
      TFun arrowKind lhs rhs ->
        [lhs, rhs]
          <> [TFun arrowKind lhs' rhs | lhs' <- shrinkType lhs]
          <> [TFun arrowKind lhs rhs' | rhs' <- shrinkType rhs]
      TInfix lhs op promoted rhs ->
        [lhs, rhs]
          <> [TInfix lhs' op promoted rhs | lhs' <- shrinkType lhs]
          <> [TInfix lhs op promoted rhs' | rhs' <- shrinkType rhs]
      TTuple tupleFlavor _ elems ->
        shrinkTypeTupleElems tupleFlavor elems
      TList _ elems ->
        [TList Unpromoted elems' | elems' <- shrinkList shrinkType elems, not (null elems')]
          <> elems
      TParen inner ->
        [inner]
          <> [TParen inner' | inner' <- shrinkType inner]
      TKindSig ty' kind ->
        [ty', kind]
          <> [TKindSig ty'' kind | ty'' <- shrinkType ty']
          <> [TKindSig ty' kind' | kind' <- shrinkType kind]
      TUnboxedSum elems ->
        [TUnboxedSum elems' | elems' <- shrinkList shrinkType elems, length elems' >= 2]
      TContext constraints inner ->
        [inner]
          <> [TContext constraints' inner | constraints' <- shrinkList shrinkType constraints, not (null constraints')]
          <> [TContext constraints inner' | inner' <- shrinkType inner]
      TSplice e ->
        [TSplice e' | e' <- shrinkExpr e]
      TWildcard ->
        []
      TAnn _ sub -> shrinkType sub

shrinkTyVarBinders :: [TyVarBinder] -> [[TyVarBinder]]
shrinkTyVarBinders binders =
  [ shrunk
  | shrunk <- shrinkList shrinkTyVarBinder binders,
    not (null shrunk)
  ]

shrinkTyVarBinder :: TyVarBinder -> [TyVarBinder]
shrinkTyVarBinder tvb =
  [tvb {tyVarBinderName = name'} | name' <- shrinkIdent (tyVarBinderName tvb)]
    <> [tvb {tyVarBinderKind = Just kind'} | Just kind <- [tyVarBinderKind tvb], kind' <- shrinkType kind]

shrinkForallTelescope :: ForallTelescope -> [ForallTelescope]
shrinkForallTelescope telescope =
  [ telescope {forallTelescopeBinders = binders'}
  | binders' <- shrinkTyVarBinders (forallTelescopeBinders telescope)
  ]
    <> [ telescope {forallTelescopeVisibility = ForallInvisible}
       | forallTelescopeVisibility telescope == ForallVisible
       ]

shrinkTypeTupleElems :: TupleFlavor -> [Type] -> [Type]
shrinkTypeTupleElems tupleFlavor elems =
  elems
    <> [TTuple Boxed Unpromoted elems | tupleFlavor == Unboxed, length elems /= 1]
    <> [ candidate
       | shrunk <- shrinkList shrinkType elems,
         candidate <- case shrunk of
           [] -> [TTuple tupleFlavor Unpromoted []]
           [_] | tupleFlavor == Boxed -> []
           _ -> [TTuple tupleFlavor Unpromoted shrunk]
       ]

shrinkImplicitParamName :: Text -> [Text]
shrinkImplicitParamName name =
  case T.stripPrefix "?" name of
    Nothing -> []
    Just inner -> ["?" <> candidate | candidate <- shrinkIdent inner]