{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}

module Ormolu.Config.TH
  ( allNothing,
    unpackFieldsWithSuffix,

    -- * BijectiveMap
    BijectiveMap,
    mkBijectiveMap,
    parseTextWith,
    showTextWith,
    showAllValues,
  )
where

import Control.Monad (forM, (>=>))
import Data.List (intercalate, nub)
import Language.Haskell.TH
import Language.Haskell.TH.Syntax (lift)
import Text.Printf (printf)

allNothing :: Name -> Q Exp
allNothing :: Name -> Q Exp
allNothing Name
name = do
  Type
ty <- Name -> Q Type
reifyType Name
name
  (Q Exp -> Q Exp -> Q Exp) -> Q Exp -> [Q Exp] -> Q Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Q Exp -> Q Exp -> Q Exp
appE (Name -> Q Exp
conE Name
name) ([Q Exp] -> Q Exp) -> [Q Exp] -> Q Exp
forall a b. (a -> b) -> a -> b
$
    Int -> Q Exp -> [Q Exp]
forall a. Int -> a -> [a]
replicate (Type -> Int
getArity Type
ty) [|Nothing|]

unpackFieldsWithSuffix :: Name -> String -> Q Pat
unpackFieldsWithSuffix :: Name -> String -> Q Pat
unpackFieldsWithSuffix Name
name String
suffix = do
  Name
typeForCon <-
    Name -> Q Info
reify Name
name Q Info -> (Info -> Q Name) -> Q Name
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      DataConI Name
_ Type
_ Name
typeForCon -> Name -> Q Name
forall (m :: * -> *) a. Monad m => a -> m a
return Name
typeForCon
      Info
info -> String -> Q Name
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Name) -> String -> Q Name
forall a b. (a -> b) -> a -> b
$ String
"allNothing requires the Name of a data constructor, got: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Info -> String
forall a. Show a => a -> String
show Info
info

  [Con]
allConsInType <-
    Name -> Q (Either Info [Con])
getAllConstructors Name
typeForCon
      Q (Either Info [Con]) -> (Either Info [Con] -> Q [Con]) -> Q [Con]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Info -> Q [Con])
-> ([Con] -> Q [Con]) -> Either Info [Con] -> Q [Con]
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Q [Con]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q [Con]) -> (Info -> String) -> Info -> Q [Con]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String -> String
forall r. PrintfType r => String -> r
printf String
"Unexpected parent of data constructor: %s" (String -> String) -> (Info -> String) -> Info -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Info -> String
forall a. Show a => a -> String
show) [Con] -> Q [Con]
forall (m :: * -> *) a. Monad m => a -> m a
return

  [Name]
fields <-
    case (Con -> Bool) -> [Con] -> [Con]
forall a. (a -> Bool) -> [a] -> [a]
filter (Name -> [Name] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem Name
name ([Name] -> Bool) -> (Con -> [Name]) -> Con -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> [Name]
getConstructorNames) [Con]
allConsInType of
      [Con
con] | Just [Name]
fields <- Con -> Maybe [Name]
conFieldNames Con
con -> [Name] -> Q [Name]
forall (m :: * -> *) a. Monad m => a -> m a
return [Name]
fields
      [Con]
_ -> String -> Q [Name]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q [Name]) -> String -> Q [Name]
forall a b. (a -> b) -> a -> b
$ String
"Could not find unique record constructor in: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [Con] -> String
forall a. Show a => a -> String
show [Con]
allConsInType

  Name -> [Q Pat] -> Q Pat
conP Name
name ([Q Pat] -> Q Pat) -> [Q Pat] -> Q Pat
forall a b. (a -> b) -> a -> b
$ (Name -> Q Pat) -> [Name] -> [Q Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Q Pat
varP (Name -> Q Pat) -> (Name -> Name) -> Name -> Q Pat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Name
mkName (String -> Name) -> (Name -> String) -> Name -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
suffix) (String -> String) -> (Name -> String) -> Name -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase) [Name]
fields
  where
    conFieldNames :: Con -> Maybe [Name]
conFieldNames = \case
      NormalC {} -> Maybe [Name]
forall a. Maybe a
Nothing
      RecC Name
_ [VarBangType]
tys -> [Name] -> Maybe [Name]
forall a. a -> Maybe a
Just ([Name] -> Maybe [Name]) -> [Name] -> Maybe [Name]
forall a b. (a -> b) -> a -> b
$ (VarBangType -> Name) -> [VarBangType] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map VarBangType -> Name
forall a b c. (a, b, c) -> a
fst3 [VarBangType]
tys
      InfixC {} -> Maybe [Name]
forall a. Maybe a
Nothing
      ForallC {} -> Maybe [Name]
forall a. Maybe a
Nothing
      GadtC {} -> Maybe [Name]
forall a. Maybe a
Nothing
      RecGadtC [Name]
_ [VarBangType]
tys Type
_ -> [Name] -> Maybe [Name]
forall a. a -> Maybe a
Just ([Name] -> Maybe [Name]) -> [Name] -> Maybe [Name]
forall a b. (a -> b) -> a -> b
$ (VarBangType -> Name) -> [VarBangType] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map VarBangType -> Name
forall a b c. (a, b, c) -> a
fst3 [VarBangType]
tys
    fst3 :: (a, b, c) -> a
fst3 (a
x, b
_, c
_) = a
x

data BijectiveMap a = BijectiveMap
  { BijectiveMap a -> String -> Either String a
parseTextWith :: String -> Either String a,
    BijectiveMap a -> a -> String
showTextWith :: a -> String,
    BijectiveMap a -> [String]
getAllOptions :: [String]
  }

showAllValues :: BijectiveMap a -> String
showAllValues :: BijectiveMap a -> String
showAllValues = [String] -> String
uncommas ([String] -> String)
-> (BijectiveMap a -> [String]) -> BijectiveMap a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> String) -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map String -> String
forall a. Show a => a -> String
show ([String] -> [String])
-> (BijectiveMap a -> [String]) -> BijectiveMap a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BijectiveMap a -> [String]
forall a. BijectiveMap a -> [String]
getAllOptions

-- | Generate a `BijectiveMap a` value with the given map.
--
-- Checks the following:
--   * all Names in given list refer to a constructor of type `a`
--   * all Names in given list refer to a 0-arity constructor
--   * all constructors in type `a` are accounted for.
mkBijectiveMap :: [(Name, String)] -> Q Exp
mkBijectiveMap :: [(Name, String)] -> Q Exp
mkBijectiveMap [(Name, String)]
mapping = do
  let ([Name]
conNames, [String]
allOptions) = [(Name, String)] -> ([Name], [String])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Name, String)]
mapping

  -- check all names refer to constructors
  ([Type]
conTypes, [Name]
conParents) <-
    ([(Type, Name)] -> ([Type], [Name]))
-> Q [(Type, Name)] -> Q ([Type], [Name])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(Type, Name)] -> ([Type], [Name])
forall a b. [(a, b)] -> ([a], [b])
unzip (Q [(Type, Name)] -> Q ([Type], [Name]))
-> ((Name -> Q (Type, Name)) -> Q [(Type, Name)])
-> (Name -> Q (Type, Name))
-> Q ([Type], [Name])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Name] -> (Name -> Q (Type, Name)) -> Q [(Type, Name)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Name]
conNames ((Name -> Q (Type, Name)) -> Q ([Type], [Name]))
-> (Name -> Q (Type, Name)) -> Q ([Type], [Name])
forall a b. (a -> b) -> a -> b
$ \Name
name ->
      Name -> Q Info
reify Name
name Q Info -> (Info -> Q (Type, Name)) -> Q (Type, Name)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        DataConI Name
_ Type
ty Name
parent -> (Type, Name) -> Q (Type, Name)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
ty, Name
parent)
        Info
info ->
          String -> Q (Type, Name)
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q (Type, Name)) -> String -> Q (Type, Name)
forall a b. (a -> b) -> a -> b
$
            String -> String -> String -> String
forall r. PrintfType r => String -> r
printf
              String
"mkBijectiveMap requires all Names refer to data constructors, got %s: %s"
              (Name -> String
forall a. Show a => a -> String
show Name
name)
              (Info -> String
forall a. Show a => a -> String
show Info
info)

  -- check that all constructors are in same type
  Name
parent <-
    case [Name] -> [Name]
forall a. Eq a => [a] -> [a]
nub [Name]
conParents of
      [Name
parent] -> Name -> Q Name
forall (m :: * -> *) a. Monad m => a -> m a
return Name
parent
      [Name]
parents -> String -> Q Name
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Name) -> String -> Q Name
forall a b. (a -> b) -> a -> b
$ String
"mkBijectiveMap requires all Names refer to data constructors in the same type, got: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [Name] -> String
forall a. Show a => a -> String
show [Name]
parents

  -- check that all constructors are 0-arity
  case (Type -> Bool) -> [Type] -> [Type]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (Int -> Bool) -> (Type -> Int) -> Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Int
getArity) [Type]
conTypes of
    [] -> () -> Q ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    [Type]
_ -> String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"mkBijectiveMap requires all constructors have 0-arity"

  -- check that all constructors are represented
  [Con]
allConsInType <-
    Name -> Q (Either Info [Con])
getAllConstructors Name
parent
      Q (Either Info [Con]) -> (Either Info [Con] -> Q [Con]) -> Q [Con]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Info -> Q [Con])
-> ([Con] -> Q [Con]) -> Either Info [Con] -> Q [Con]
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Q [Con]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q [Con]) -> (Info -> String) -> Info -> Q [Con]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String -> String
forall r. PrintfType r => String -> r
printf String
"Unexpected parent of data constructors: %s" (String -> String) -> (Info -> String) -> Info -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Info -> String
forall a. Show a => a -> String
show) [Con] -> Q [Con]
forall (m :: * -> *) a. Monad m => a -> m a
return
  case (Name -> Bool) -> [Name] -> [Name]
forall a. (a -> Bool) -> [a] -> [a]
filter (Name -> [Name] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Name]
conNames) ((Con -> [Name]) -> [Con] -> [Name]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Con -> [Name]
getConstructorNames [Con]
allConsInType) of
    [] -> () -> Q ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    [Name]
missing -> String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q ()) -> String -> Q ()
forall a b. (a -> b) -> a -> b
$ String
"Missing constructors: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [Name] -> String
forall a. Show a => a -> String
show [Name]
missing

  Name
unknown <- String -> Q Name
newName String
"unknown"
  let parser :: Q Exp
parser =
        [MatchQ] -> Q Exp
lamCaseE ([MatchQ] -> Q Exp)
-> ([[MatchQ]] -> [MatchQ]) -> [[MatchQ]] -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[MatchQ]] -> [MatchQ]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[MatchQ]] -> Q Exp) -> [[MatchQ]] -> Q Exp
forall a b. (a -> b) -> a -> b
$
          [ (((Name, String) -> MatchQ) -> [(Name, String)] -> [MatchQ])
-> [(Name, String)] -> ((Name, String) -> MatchQ) -> [MatchQ]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Name, String) -> MatchQ) -> [(Name, String)] -> [MatchQ]
forall a b. (a -> b) -> [a] -> [b]
map [(Name, String)]
mapping (((Name, String) -> MatchQ) -> [MatchQ])
-> ((Name, String) -> MatchQ) -> [MatchQ]
forall a b. (a -> b) -> a -> b
$ \(Name
name, String
option) ->
              Q Pat -> BodyQ -> [DecQ] -> MatchQ
match
                (Lit -> Q Pat
litP (Lit -> Q Pat) -> Lit -> Q Pat
forall a b. (a -> b) -> a -> b
$ String -> Lit
stringL String
option)
                (Q Exp -> BodyQ
normalB [|Right $(conE name)|])
                [],
            [ Q Pat -> BodyQ -> [DecQ] -> MatchQ
match
                (Name -> Q Pat
varP Name
unknown)
                ( Q Exp -> BodyQ
normalB
                    [|
                      Left . unlines $
                        [ "unknown value: " <> show $(varE unknown),
                          "Valid values are: " <> $(lift $ uncommas $ map show allOptions)
                        ]
                      |]
                )
                []
            ]
          ]
      shower :: Q Exp
shower =
        [MatchQ] -> Q Exp
lamCaseE ([MatchQ] -> Q Exp) -> [MatchQ] -> Q Exp
forall a b. (a -> b) -> a -> b
$
          (((Name, String) -> MatchQ) -> [(Name, String)] -> [MatchQ])
-> [(Name, String)] -> ((Name, String) -> MatchQ) -> [MatchQ]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Name, String) -> MatchQ) -> [(Name, String)] -> [MatchQ]
forall a b. (a -> b) -> [a] -> [b]
map [(Name, String)]
mapping (((Name, String) -> MatchQ) -> [MatchQ])
-> ((Name, String) -> MatchQ) -> [MatchQ]
forall a b. (a -> b) -> a -> b
$ \(Name
name, String
option) ->
            Q Pat -> BodyQ -> [DecQ] -> MatchQ
match (Name -> [Q Pat] -> Q Pat
conP Name
name []) (Q Exp -> BodyQ
normalB (Q Exp -> BodyQ) -> Q Exp -> BodyQ
forall a b. (a -> b) -> a -> b
$ String -> Q Exp
forall t. Lift t => t -> Q Exp
lift String
option) []

  [|
    BijectiveMap
      { parseTextWith = $parser,
        showTextWith = $shower,
        getAllOptions = $(lift allOptions)
      }
    |]

----------------------------------------------------------------------------
-- Helpers

{- FOURMOLU_DISABLE -}
{- https://github.com/fourmolu/fourmolu#limitations -}
getArity :: Type -> Int
getArity :: Type -> Int
getArity = \case
  ForallT [TyVarBndr]
_ [Type]
_ Type
ty -> Type -> Int
getArity Type
ty
  AppT (AppT Type
ArrowT Type
_) Type
ty -> Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Type -> Int
getArity Type
ty
#if MIN_VERSION_template_haskell(2,17,0)
  AppT (AppT (AppT MulArrowT _) _) ty -> 1 + getArity ty
#endif
  Type
_ -> Int
0
{- FOURMOLU_ENABLE -}

getAllConstructors :: Name -> Q (Either Info [Con])
getAllConstructors :: Name -> Q (Either Info [Con])
getAllConstructors =
  Name -> Q Info
reify (Name -> Q Info)
-> (Info -> Q (Either Info [Con])) -> Name -> Q (Either Info [Con])
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> \case
    TyConI (DataD [Type]
_ Name
_ [TyVarBndr]
_ Maybe Type
_ [Con]
cons [DerivClause]
_) -> Either Info [Con] -> Q (Either Info [Con])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Info [Con] -> Q (Either Info [Con]))
-> Either Info [Con] -> Q (Either Info [Con])
forall a b. (a -> b) -> a -> b
$ [Con] -> Either Info [Con]
forall a b. b -> Either a b
Right [Con]
cons
    Info
info -> Either Info [Con] -> Q (Either Info [Con])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Info [Con] -> Q (Either Info [Con]))
-> Either Info [Con] -> Q (Either Info [Con])
forall a b. (a -> b) -> a -> b
$ Info -> Either Info [Con]
forall a b. a -> Either a b
Left Info
info

-- Could return multiple names for GADTs like 'A, B :: Foo'
getConstructorNames :: Con -> [Name]
getConstructorNames :: Con -> [Name]
getConstructorNames = \case
  NormalC Name
n [BangType]
_ -> [Name
n]
  RecC Name
n [VarBangType]
_ -> [Name
n]
  InfixC BangType
_ Name
n BangType
_ -> [Name
n]
  ForallC [TyVarBndr]
_ [Type]
_ Con
c -> Con -> [Name]
getConstructorNames Con
c
  GadtC [Name]
ns [BangType]
_ Type
_ -> [Name]
ns
  RecGadtC [Name]
ns [VarBangType]
_ Type
_ -> [Name]
ns

uncommas :: [String] -> String
uncommas :: [String] -> String
uncommas [] = String
""
uncommas [String
s] = String
s
uncommas [String
s0, String
s1] = String
s0 String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" or " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
s1
uncommas [String]
ss =
  let pre :: [String]
pre = [String] -> [String]
forall a. [a] -> [a]
init [String]
ss
      end :: String
end = [String] -> String
forall a. [a] -> a
last [String]
ss
   in String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " [String]
pre String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"or " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
end