{-# LANGUAGE CPP                 #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGuaGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns        #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGuaGE DeriveGeneric       #-}
{-# LANGuaGE FlexibleContexts    #-}
-- | Formatting type declarations and class instances for inferred types. 
module Data.Aeson.AutoType.Split(
  splitTypeByLabel, unificationCandidates,
  unifyCandidates, toposort
) where

import           Control.Arrow             ((&&&))
import           Control.Applicative       ((<$>), (<*>))
import           Control.Lens.TH
import           Control.Lens
import           Control.Monad             (forM)
import           Control.Exception(assert)
import qualified Data.HashMap.Strict        as Map
import           Data.Monoid
import qualified Data.Set                   as Set
import qualified Data.Text                  as Text
import           Data.Text                 (Text)
import           Data.Set                  (Set )
import           Data.List                 (foldl1')
import           Data.Char                 (isAlpha, isDigit)
import           Control.Monad.State.Class
import           Control.Monad.State.Strict(State, runState)
import qualified Data.Graph          as Graph
import           GHC.Generics              (Generic)

import           Data.Aeson.AutoType.Type
import           Data.Aeson.AutoType.Extract
import           Data.Aeson.AutoType.Util  ()

--import           Debug.Trace -- DEBUG
trace :: p -> p -> p
trace _ x :: p
x = p
x

fst3 ::  (t, t1, t2) -> t
fst3 :: (t, t1, t2) -> t
fst3 (a :: t
a, _b :: t1
_b, _c :: t2
_c) = t
a

type Map k v = Map.HashMap k v 

-- | Explanatory type alias for making declarations
-- First element of the triple is original JSON identifier,
-- second element of the triple is the mapped identifier name in Haskell.
-- third element of the triple shows the type in a formatted way
type MappedKey = (Text, Text, Text, Bool)

-- * Splitting object types by label for unification.
type TypeTree    = Map Text [Type]

type TypeTreeM a = State TypeTree a

addType :: Text -> Type -> TypeTreeM ()
addType :: Text -> Type -> TypeTreeM ()
addType label :: Text
label typ :: Type
typ = (HashMap Text [Type] -> HashMap Text [Type]) -> TypeTreeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((HashMap Text [Type] -> HashMap Text [Type]) -> TypeTreeM ())
-> (HashMap Text [Type] -> HashMap Text [Type]) -> TypeTreeM ()
forall a b. (a -> b) -> a -> b
$ ([Type] -> [Type] -> [Type])
-> Text -> [Type] -> HashMap Text [Type] -> HashMap Text [Type]
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
Map.insertWith [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
(++) Text
label [Type
typ]

splitTypeByLabel' :: Text -> Type -> TypeTreeM Type
splitTypeByLabel' :: Text -> Type -> TypeTreeM Type
splitTypeByLabel' _  TString   = Type -> TypeTreeM Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
TString
splitTypeByLabel' _  TInt      = Type -> TypeTreeM Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
TInt
splitTypeByLabel' _  TDouble   = Type -> TypeTreeM Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
TDouble
splitTypeByLabel' _  TBool     = Type -> TypeTreeM Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
TBool
splitTypeByLabel' _  TNull     = Type -> TypeTreeM Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
TNull
splitTypeByLabel' _ (TLabel r :: Text
r) = Bool -> TypeTreeM Type -> TypeTreeM Type
forall a. (?callStack::CallStack) => Bool -> a -> a
assert Bool
False (TypeTreeM Type -> TypeTreeM Type)
-> TypeTreeM Type -> TypeTreeM Type
forall a b. (a -> b) -> a -> b
$ Type -> TypeTreeM Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> TypeTreeM Type) -> Type -> TypeTreeM Type
forall a b. (a -> b) -> a -> b
$ Text -> Type
TLabel Text
r -- unnecessary?
splitTypeByLabel' l :: Text
l (TUnion u :: Set Type
u) = do [Type]
m <- (Type -> TypeTreeM Type)
-> [Type] -> StateT (HashMap Text [Type]) Identity [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Text -> Type -> TypeTreeM Type
splitTypeByLabel' Text
l) ([Type] -> StateT (HashMap Text [Type]) Identity [Type])
-> [Type] -> StateT (HashMap Text [Type]) Identity [Type]
forall a b. (a -> b) -> a -> b
$ Set Type -> [Type]
forall a. Set a -> [a]
Set.toList Set Type
u
                                    Type -> TypeTreeM Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> TypeTreeM Type) -> Type -> TypeTreeM Type
forall a b. (a -> b) -> a -> b
$! Set Type -> Type
TUnion (Set Type -> Type) -> Set Type -> Type
forall a b. (a -> b) -> a -> b
$! [Type] -> Set Type
forall a. Ord a => [a] -> Set a
Set.fromList [Type]
m
splitTypeByLabel' l :: Text
l (TArray a :: Type
a) = do Type
m <- Text -> Type -> TypeTreeM Type
splitTypeByLabel' (Text
l Text -> Text -> Text
`Text.append` "Elt") Type
a
                                    Type -> TypeTreeM Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> TypeTreeM Type) -> Type -> TypeTreeM Type
forall a b. (a -> b) -> a -> b
$! Type -> Type
TArray Type
m
splitTypeByLabel' l :: Text
l (TObj   o :: Dict
o) = do [(Text, Type)]
kvs <- [(Text, Type)]
-> ((Text, Type)
    -> StateT (HashMap Text [Type]) Identity (Text, Type))
-> StateT (HashMap Text [Type]) Identity [(Text, Type)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (HashMap Text Type -> [(Text, Type)]
forall k v. HashMap k v -> [(k, v)]
Map.toList (HashMap Text Type -> [(Text, Type)])
-> HashMap Text Type -> [(Text, Type)]
forall a b. (a -> b) -> a -> b
$ Dict -> HashMap Text Type
unDict Dict
o) (((Text, Type)
  -> StateT (HashMap Text [Type]) Identity (Text, Type))
 -> StateT (HashMap Text [Type]) Identity [(Text, Type)])
-> ((Text, Type)
    -> StateT (HashMap Text [Type]) Identity (Text, Type))
-> StateT (HashMap Text [Type]) Identity [(Text, Type)]
forall a b. (a -> b) -> a -> b
$ \(k :: Text
k, v :: Type
v) -> do
                                       Type
component <- Text -> Type -> TypeTreeM Type
splitTypeByLabel' Text
k Type
v
                                       (Text, Type) -> StateT (HashMap Text [Type]) Identity (Text, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (Text
k, Type
component)
                                    Text -> Type -> TypeTreeM ()
addType Text
l (Dict -> Type
TObj (Dict -> Type) -> Dict -> Type
forall a b. (a -> b) -> a -> b
$ HashMap Text Type -> Dict
Dict (HashMap Text Type -> Dict) -> HashMap Text Type -> Dict
forall a b. (a -> b) -> a -> b
$ [(Text, Type)] -> HashMap Text Type
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList [(Text, Type)]
kvs)
                                    Type -> TypeTreeM Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> TypeTreeM Type) -> Type -> TypeTreeM Type
forall a b. (a -> b) -> a -> b
$! Text -> Type
TLabel Text
l

-- | Splits initial type with a given label, into a mapping of object type names and object type structures.
splitTypeByLabel :: Text -> Type -> Map Text Type
splitTypeByLabel :: Text -> Type -> HashMap Text Type
splitTypeByLabel topLabel :: Text
topLabel t :: Type
t = ([Type] -> Type) -> HashMap Text [Type] -> HashMap Text Type
forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
Map.map ((Type -> Type -> Type) -> [Type] -> Type
forall a. (a -> a -> a) -> [a] -> a
foldl1' Type -> Type -> Type
unifyTypes) HashMap Text [Type]
finalState
  where
    finalize :: Type -> TypeTreeM ()
finalize (TLabel l :: Text
l) = Bool -> TypeTreeM () -> TypeTreeM ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Text
l Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
topLabel) (TypeTreeM () -> TypeTreeM ()) -> TypeTreeM () -> TypeTreeM ()
forall a b. (a -> b) -> a -> b
$ () -> TypeTreeM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    finalize  topLevel :: Type
topLevel  = Text -> Type -> TypeTreeM ()
addType Text
topLabel Type
topLevel
    initialState :: HashMap k v
initialState    = HashMap k v
forall k v. HashMap k v
Map.empty
    (_, finalState :: HashMap Text [Type]
finalState) = TypeTreeM () -> HashMap Text [Type] -> ((), HashMap Text [Type])
forall s a. State s a -> s -> (a, s)
runState (Text -> Type -> TypeTreeM Type
splitTypeByLabel' Text
topLabel Type
t TypeTreeM Type -> (Type -> TypeTreeM ()) -> TypeTreeM ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Type -> TypeTreeM ()
finalize) HashMap Text [Type]
forall k v. HashMap k v
initialState

-- | Topological sorting of splitted types so that it is accepted declaration order.
toposort :: Map Text Type -> [(Text, Type)]  
toposort :: HashMap Text Type -> [(Text, Type)]
toposort splitted :: HashMap Text Type
splitted = (Vertex -> (Text, Type)) -> [Vertex] -> [(Text, Type)]
forall a b. (a -> b) -> [a] -> [b]
map ((Text -> Text
forall a. a -> a
id (Text -> Text) -> (Text -> Type) -> Text -> (Text, Type)
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& (HashMap Text Type
splitted HashMap Text Type -> Text -> Type
forall k v. (Eq k, Hashable k) => HashMap k v -> k -> v
Map.!)) (Text -> (Text, Type))
-> (Vertex -> Text) -> Vertex -> (Text, Type)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, Text, [Text]) -> Text
forall t t1 t2. (t, t1, t2) -> t
fst3 ((Text, Text, [Text]) -> Text)
-> (Vertex -> (Text, Text, [Text])) -> Vertex -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vertex -> (Text, Text, [Text])
graphKey) ([Vertex] -> [(Text, Type)]) -> [Vertex] -> [(Text, Type)]
forall a b. (a -> b) -> a -> b
$ Graph -> [Vertex]
Graph.topSort Graph
graph
  where
    (graph :: Graph
graph, graphKey :: Vertex -> (Text, Text, [Text])
graphKey) = [(Text, Text, [Text])] -> (Graph, Vertex -> (Text, Text, [Text]))
forall key node.
Ord key =>
[(node, key, [key])] -> (Graph, Vertex -> (node, key, [key]))
Graph.graphFromEdges' ([(Text, Text, [Text])] -> (Graph, Vertex -> (Text, Text, [Text])))
-> [(Text, Text, [Text])]
-> (Graph, Vertex -> (Text, Text, [Text]))
forall a b. (a -> b) -> a -> b
$ ((Text, Type) -> (Text, Text, [Text]))
-> [(Text, Type)] -> [(Text, Text, [Text])]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Type) -> (Text, Text, [Text])
forall b. (b, Type) -> (b, b, [Text])
makeEntry ([(Text, Type)] -> [(Text, Text, [Text])])
-> [(Text, Type)] -> [(Text, Text, [Text])]
forall a b. (a -> b) -> a -> b
$ HashMap Text Type -> [(Text, Type)]
forall k v. HashMap k v -> [(k, v)]
Map.toList HashMap Text Type
splitted
    makeEntry :: (b, Type) -> (b, b, [Text])
makeEntry (k :: b
k, v :: Type
v) = (b
k, b
k, Type -> [Text]
allLabels Type
v)

-- | Computes all type labels referenced by a given type.
allLabels :: Type -> [Text]
allLabels :: Type -> [Text]
allLabels = (Type -> [Text] -> [Text]) -> [Text] -> Type -> [Text]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Type -> [Text] -> [Text]
go []
  where
    go :: Type -> [Text] -> [Text]
go (TLabel l :: Text
l) ls :: [Text]
ls = Text
lText -> [Text] -> [Text]
forall a. a -> [a] -> [a]
:[Text]
ls
    go (TArray t :: Type
t) ls :: [Text]
ls = Type -> [Text] -> [Text]
go Type
t [Text]
ls
    go (TUnion u :: Set Type
u) ls :: [Text]
ls = (Type -> [Text] -> [Text]) -> [Text] -> Set Type -> [Text]
forall a b. (a -> b -> b) -> b -> Set a -> b
Set.foldr Type -> [Text] -> [Text]
go [Text]
ls          Set Type
u
    go (TObj   o :: Dict
o) ls :: [Text]
ls = (Type -> [Text] -> [Text]) -> [Text] -> HashMap Text Type -> [Text]
forall v a k. (v -> a -> a) -> a -> HashMap k v -> a
Map.foldr Type -> [Text] -> [Text]
go [Text]
ls (HashMap Text Type -> [Text]) -> HashMap Text Type -> [Text]
forall a b. (a -> b) -> a -> b
$ Dict -> HashMap Text Type
unDict Dict
o
    go _other :: Type
_other     ls :: [Text]
ls = [Text]
ls

-- * Finding candidates for extra unifications
-- | For a given splitted types, it returns candidates for extra
-- unifications.
unificationCandidates :: Map.HashMap t Type -> [[t]]
unificationCandidates :: HashMap t Type -> [[t]]
unificationCandidates = HashMap (Set Text) [t] -> [[t]]
forall k v. HashMap k v -> [v]
Map.elems             (HashMap (Set Text) [t] -> [[t]])
-> (HashMap t Type -> HashMap (Set Text) [t])
-> HashMap t Type
-> [[t]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                        ([t] -> Bool) -> HashMap (Set Text) [t] -> HashMap (Set Text) [t]
forall v k. (v -> Bool) -> HashMap k v -> HashMap k v
Map.filter [t] -> Bool
forall a. [a] -> Bool
candidates (HashMap (Set Text) [t] -> HashMap (Set Text) [t])
-> (HashMap t Type -> HashMap (Set Text) [t])
-> HashMap t Type
-> HashMap (Set Text) [t]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                        ([t] -> [t] -> [t]) -> [(Set Text, [t])] -> HashMap (Set Text) [t]
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> [(k, v)] -> HashMap k v
Map.fromListWith [t] -> [t] -> [t]
forall a. [a] -> [a] -> [a]
(++) ([(Set Text, [t])] -> HashMap (Set Text) [t])
-> (HashMap t Type -> [(Set Text, [t])])
-> HashMap t Type
-> HashMap (Set Text) [t]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                        ((t, Type) -> [(Set Text, [t])])
-> [(t, Type)] -> [(Set Text, [t])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (t, Type) -> [(Set Text, [t])]
forall a. (a, Type) -> [(Set Text, [a])]
entry       ([(t, Type)] -> [(Set Text, [t])])
-> (HashMap t Type -> [(t, Type)])
-> HashMap t Type
-> [(Set Text, [t])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                        HashMap t Type -> [(t, Type)]
forall k v. HashMap k v -> [(k, v)]
Map.toList
  where
    -- | Candidate entry has to have at least two candidates, so that unification makes sense
    candidates :: [a] -> Bool
candidates [ ] = Bool
False
    candidates [_] = Bool
False
    candidates _   = Bool
True
    -- | Make a candidate entry for each object type, which points from its keys to its label.
    entry :: (a, Type) -> [(Set Text, [a])]
entry (k :: a
k, TObj o :: Dict
o)                 = [([Text] -> Set Text
forall a. Ord a => [a] -> Set a
Set.fromList ([Text] -> Set Text) -> [Text] -> Set Text
forall a b. (a -> b) -> a -> b
$ HashMap Text Type -> [Text]
forall k v. HashMap k v -> [k]
Map.keys (HashMap Text Type -> [Text]) -> HashMap Text Type -> [Text]
forall a b. (a -> b) -> a -> b
$ Dict -> HashMap Text Type
unDict Dict
o, [a
k])]
    entry  _                          = [] -- ignore array elements and toplevel type if it is Array

-- | Unifies candidates on a give input list.
unifyCandidates :: [[Text]] -> Map Text Type -> Map Text Type
unifyCandidates :: [[Text]] -> HashMap Text Type -> HashMap Text Type
unifyCandidates candidates :: [[Text]]
candidates splitted :: HashMap Text Type
splitted = (Type -> Type) -> HashMap Text Type -> HashMap Text Type
forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
Map.map (Map Text Text -> Type -> Type
remapLabels Map Text Text
labelMapping) (HashMap Text Type -> HashMap Text Type)
-> HashMap Text Type -> HashMap Text Type
forall a b. (a -> b) -> a -> b
$ HashMap Text Type -> HashMap Text Type
replacements HashMap Text Type
splitted
  where
    unifiedType  :: [Text] -> Type
    unifiedType :: [Text] -> Type
unifiedType cset :: [Text]
cset      = (Type -> Type -> Type) -> [Type] -> Type
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Type -> Type -> Type
unifyTypes         ([Type] -> Type) -> [Type] -> Type
forall a b. (a -> b) -> a -> b
$ 
                            (Text -> Type) -> [Text] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (HashMap Text Type
splitted HashMap Text Type -> Text -> Type
forall k v. (Eq k, Hashable k) => HashMap k v -> k -> v
Map.!) [Text]
cset
    replace      :: [Text] -> Map Text Type -> Map Text Type
    replace :: [Text] -> HashMap Text Type -> HashMap Text Type
replace  cset :: [Text]
cset@(c :: Text
c:_) s :: HashMap Text Type
s = Text -> Type -> HashMap Text Type -> HashMap Text Type
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert Text
c ([Text] -> Type
unifiedType [Text]
cset) ((Text -> HashMap Text Type -> HashMap Text Type)
-> HashMap Text Type -> [Text] -> HashMap Text Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Text -> HashMap Text Type -> HashMap Text Type
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
Map.delete HashMap Text Type
s [Text]
cset)
    replace  []         _ = [Char] -> HashMap Text Type
forall a. (?callStack::CallStack) => [Char] -> a
error "Empty candidate set in replace"
    replacements :: Map Text Type -> Map Text Type
    replacements :: HashMap Text Type -> HashMap Text Type
replacements        s :: HashMap Text Type
s = ([Text] -> HashMap Text Type -> HashMap Text Type)
-> HashMap Text Type -> [[Text]] -> HashMap Text Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr [Text] -> HashMap Text Type -> HashMap Text Type
replace HashMap Text Type
s [[Text]]
candidates
    labelMapping :: Map Text Text
    labelMapping :: Map Text Text
labelMapping          = [(Text, Text)] -> Map Text Text
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(Text, Text)] -> Map Text Text)
-> [(Text, Text)] -> Map Text Text
forall a b. (a -> b) -> a -> b
$ ([Text] -> [(Text, Text)]) -> [[Text]] -> [(Text, Text)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [Text] -> [(Text, Text)]
forall b. [b] -> [(b, b)]
mapEntry [[Text]]
candidates
    mapEntry :: [b] -> [(b, b)]
mapEntry cset :: [b]
cset@(c :: b
c:_)   = [(b
x, b
c) | b
x <- [b]
cset]
    mapEntry []           = [Char] -> [(b, b)]
forall a. (?callStack::CallStack) => [Char] -> a
error "Empty candidate set in mapEntry"

-- | Remaps type labels according to a `Map`.
remapLabels :: Map Text Text -> Type -> Type
remapLabels :: Map Text Text -> Type -> Type
remapLabels ls :: Map Text Text
ls (TObj   o :: Dict
o) = Dict -> Type
TObj   (Dict -> Type) -> Dict -> Type
forall a b. (a -> b) -> a -> b
$ HashMap Text Type -> Dict
Dict (HashMap Text Type -> Dict) -> HashMap Text Type -> Dict
forall a b. (a -> b) -> a -> b
$ (Type -> Type) -> HashMap Text Type -> HashMap Text Type
forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
Map.map (Map Text Text -> Type -> Type
remapLabels Map Text Text
ls) (HashMap Text Type -> HashMap Text Type)
-> HashMap Text Type -> HashMap Text Type
forall a b. (a -> b) -> a -> b
$ Dict -> HashMap Text Type
unDict Dict
o
remapLabels ls :: Map Text Text
ls (TArray t :: Type
t) = Type -> Type
TArray (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$                 Map Text Text -> Type -> Type
remapLabels Map Text Text
ls  Type
t
remapLabels ls :: Map Text Text
ls (TUnion u :: Set Type
u) = Set Type -> Type
TUnion (Set Type -> Type) -> Set Type -> Type
forall a b. (a -> b) -> a -> b
$        (Type -> Type) -> Set Type -> Set Type
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map (Map Text Text -> Type -> Type
remapLabels Map Text Text
ls) Set Type
u
remapLabels ls :: Map Text Text
ls (TLabel l :: Text
l) = Text -> Type
TLabel (Text -> Type) -> Text -> Type
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Map Text Text -> Text
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
Map.lookupDefault Text
l Text
l Map Text Text
ls
remapLabels _  other :: Type
other      = Type
other