content: {-# LANGUAGE BangPatterns       #-}
  {-# LANGUAGE CPP                #-}
  {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE PatternSynonyms    #-}
  {-# LANGUAGE TemplateHaskell    #-}
  {-# LANGUAGE TypeApplications   #-}
  {-# LANGUAGE ViewPatterns       #-}
  
  module PlutusTx.AsData (asData, asDataFor) where
  
  import Control.Lens (ifor)
  import Control.Monad (unless)
  import Data.Traversable (for)
  
  import Language.Haskell.TH qualified as TH
  import Language.Haskell.TH.Datatype qualified as TH
  import Language.Haskell.TH.Datatype.TyVarBndr qualified as TH
  
  import PlutusTx.Builtins qualified as Builtins
  import PlutusTx.IsData.Class (ToData, UnsafeFromData)
  import PlutusTx.IsData.TH (AsDataProdType (..), mkConstrCreateExpr, mkUnsafeConstrMatchPattern)
  
  import Prelude
  
  {-| 'asData' takes a datatype declaration and "backs it" by 'BuiltinData'. It does
  this by replacing the datatype with a newtype around 'BuiltinData', and providing
  pattern synonyms that match the behaviour of the original datatype.
  
  Since 'BuiltinData' can only contain 'BuiltinData', the pattern synonyms
  encode and decode all the fields using 'toBuiltinData' and 'unsafeFromBuiltinData'.
  That means that these operations are called on every pattern match or construction.
  This *can* be very efficient if, for example, the datatypes for the fields have
  also been defined with 'asData', and so have trivial conversions to/from 'BuiltinData'
  (or have very cheap conversions, like 'Integer' and 'ByteString').
  But it can be very expensive otherwise, so take care.
  
  Deriving clauses are transferred from the quoted declaration to the generated newtype
  declaration. Note that you may therefore need to do strange things like use
  @deriving newtype@ on a data declaration.
  
  __Example__:
  
  @
    $(asData [d|
        data Example a = Ex1 Integer | Ex2 a a
          deriving newtype (Eq)
      |])
  @
  
  becomes
  
  @
    newtype Example a = Example BuiltinData
      deriving newtype (Eq)
  
    pattern Ex1 :: (ToData a, UnsafeFromData a) => Integer -> Example a
    pattern Ex1 i <- Example (unsafeDataAsConstr -> ((==) 0 -> True, [unsafeFromBuiltinData -> i]))
      where Ex1 i = Example (mkConstr 0 [toBuiltinData i])
  
    pattern Ex2 :: (ToData a, UnsafeFromData a) => a -> a -> Example a
    pattern Ex2 a1 a2 <- Example (unsafeDataAsConstr -> ((==) 1 -> True,
      [ unsafeFromBuiltinData -> a1
      , unsafeFromBuiltinData -> a2
      ]))
      where Ex2 a1 a2 = Example (mkConstr 1 [toBuiltinData a1, toBuiltinData a2])
  
    {\-# COMPLETE Ex1, Ex2 #-\}
  @
  -}
  asData :: TH.Q [TH.Dec] -> TH.Q [TH.Dec]
  asData decQ = do
    decs <- decQ
    outputDecs <- for decs asDataFor
    pure $ concat outputDecs
  
  {- FOURMOLU_DISABLE -} -- Fourmolu gets upset on the segment below. This will just turn off fourmolu entirely.
  asDataFor :: TH.Dec -> TH.Q [TH.Dec]
  asDataFor dec = do
    -- th-abstraction doesn't include deriving clauses, so we have to handle that here
    let derivs = case dec of
          TH.DataD _ _ _ _ _ deriv -> deriv
          _                        -> []
  
    di@(
      TH.DatatypeInfo
        { TH.datatypeVariant = dVariant
        , TH.datatypeCons    = cons
        , TH.datatypeName    = name
        , TH.datatypeVars    = tTypeVars
        }
      ) <- TH.normalizeDec dec
  
    -- Other stuff is data families and so on
    unless (dVariant == TH.Datatype) $
      fail $ "asData: can't handle datatype variant " ++ show dVariant
    -- a fresh name for the new datatype, but same lexically as the old one
    cname <- TH.newName (show name)
    -- The newtype declaration
    let ntD =
          let con = TH.NormalC cname
                [ ( TH.Bang TH.NoSourceUnpackedness TH.NoSourceStrictness
                  , TH.ConT ''Builtins.BuiltinData
                  )
                ]
          in TH.NewtypeD [] name
  #if MIN_VERSION_template_haskell(2,21,0)
              (TH.changeTVFlags TH.BndrReq tTypeVars)
  #else
              tTypeVars
  #endif
              Nothing con derivs
  
    let isProductType =
          if length cons == 1
            then IsAsDataProdType
            else IsNotAsDataProdType
    -- The pattern synonyms, one for each constructor
    pats <- ifor cons $
      \conIx TH.ConstructorInfo
        { TH.constructorName = conName
        , TH.constructorFields = fields
        , TH.constructorVariant = cVariant
        } -> do
      -- If we have a record constructor, we need to reuse the names for the
      -- matching part of the pattern synonym
      fieldNames <- case cVariant of
        TH.RecordConstructor names -> pure names
        -- otherwise whatever
        _                          -> ifor fields $ \fieldIx _ -> TH.newName $ "arg" ++ show fieldIx
      createFieldNames <- for fieldNames (TH.newName . show)
      patSynArgs <- case cVariant of
        TH.NormalConstructor   -> pure $ TH.prefixPatSyn fieldNames
        TH.RecordConstructor _ -> pure $ TH.recordPatSyn fieldNames
        TH.InfixConstructor    -> case fieldNames of
          [f1,f2] -> pure $ TH.infixPatSyn f1 f2
          _       -> fail "asData: infix data constructor with other than two fields"
      let
        pat = TH.conP cname [mkUnsafeConstrMatchPattern isProductType (fromIntegral conIx) fieldNames]
  
        createExpr = [|$(TH.conE cname) $(mkConstrCreateExpr (fromIntegral conIx) createFieldNames) |]
        clause = TH.clause (fmap TH.varP createFieldNames) (TH.normalB createExpr) []
        patSynD = TH.patSynD conName patSynArgs (TH.explBidir [clause]) pat
        dataConstraints t = [TH.ConT ''ToData `TH.AppT` t, TH.ConT ''UnsafeFromData `TH.AppT` t]
        -- Try and be a little clever and only add constraints on the variables used in the arguments
        varsInArgs = TH.freeVariablesWellScoped fields
        ctxForArgs = concatMap (dataConstraints . TH.VarT . TH.tvName) varsInArgs
        conTy = foldr (\ty acc -> TH.ArrowT `TH.AppT` ty `TH.AppT` acc) (TH.datatypeType di) fields
        allFreeVars = TH.freeVariablesWellScoped [conTy]
        fullTy = TH.ForallT (TH.changeTVFlags TH.SpecifiedSpec allFreeVars) ctxForArgs conTy
        patSynSigD = pure $ TH.PatSynSigD conName fullTy
  
      sequence [patSynSigD, patSynD]
    -- A complete pragma, to top it off
    let compl = TH.PragmaD (TH.CompleteP (fmap TH.constructorName cons) Nothing)
    pure $ ntD : compl : concat pats
  {- FOURMOLU_ENABLE -}
  