{-# LANGUAGE NamedFieldPuns    #-}
{-# LANGUAGE OverloadedStrings #-}
module Data.Avro.Deriving.NormSchema
where

import           Control.Monad.State.Strict
import           Data.Avro.Schema.Schema
import qualified Data.Foldable              as Foldable
import qualified Data.List                  as L
import           Data.List.NonEmpty         (NonEmpty ((:|)))
import qualified Data.Map.Strict            as M
import           Data.Maybe                 (catMaybes, fromMaybe)
import           Data.Semigroup             ((<>))
import qualified Data.Set                   as S
import           Data.Text                  (Text)
import qualified Data.Text                  as T
import qualified Data.Vector                as V

-- | Extracts all the records from the schema (flattens the schema)
-- Named types get resolved when needed to include at least one "inlined"
-- schema in each record and to make each record self-contained.
-- Note: Namespaces are not really supported in this version. All the
-- namespaces (including inlined into full names) will be ignored
-- during names resolution.
extractDerivables :: Schema -> [Schema]
extractDerivables :: Schema -> [Schema]
extractDerivables Schema
s = (State (Map TypeName Schema) Schema
 -> Map TypeName Schema -> Schema)
-> Map TypeName Schema
-> State (Map TypeName Schema) Schema
-> Schema
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (Map TypeName Schema) Schema -> Map TypeName Schema -> Schema
forall s a. State s a -> s -> a
evalState Map TypeName Schema
state (State (Map TypeName Schema) Schema -> Schema)
-> ((TypeName, Schema) -> State (Map TypeName Schema) Schema)
-> (TypeName, Schema)
-> Schema
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Schema -> State (Map TypeName Schema) Schema
normSchema (Schema -> State (Map TypeName Schema) Schema)
-> ((TypeName, Schema) -> Schema)
-> (TypeName, Schema)
-> State (Map TypeName Schema) Schema
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeName, Schema) -> Schema
forall a b. (a, b) -> b
snd ((TypeName, Schema) -> Schema) -> [(TypeName, Schema)] -> [Schema]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(TypeName, Schema)]
rawRecs
  where
    rawRecs :: [(TypeName, Schema)]
rawRecs = Schema -> [(TypeName, Schema)]
getTypes Schema
s
    state :: Map TypeName Schema
state = [(TypeName, Schema)] -> Map TypeName Schema
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(TypeName, Schema)]
rawRecs

getTypes :: Schema -> [(TypeName, Schema)]
getTypes :: Schema -> [(TypeName, Schema)]
getTypes Schema
rec = case Schema
rec of
  r :: Schema
r@Record{TypeName
name :: Schema -> TypeName
name :: TypeName
name, [Field]
fields :: Schema -> [Field]
fields :: [Field]
fields} -> (TypeName
name,Schema
r) (TypeName, Schema) -> [(TypeName, Schema)] -> [(TypeName, Schema)]
forall a. a -> [a] -> [a]
: ([Field]
fields [Field] -> (Field -> [(TypeName, Schema)]) -> [(TypeName, Schema)]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Schema -> [(TypeName, Schema)]
getTypes (Schema -> [(TypeName, Schema)])
-> (Field -> Schema) -> Field -> [(TypeName, Schema)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field -> Schema
fldType))
  Array Schema
t                -> Schema -> [(TypeName, Schema)]
getTypes Schema
t
  Union Vector Schema
ts               -> (Schema -> [(TypeName, Schema)])
-> [Schema] -> [(TypeName, Schema)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Schema -> [(TypeName, Schema)]
getTypes (Vector Schema -> [Schema]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList Vector Schema
ts)
  Map Schema
t                  -> Schema -> [(TypeName, Schema)]
getTypes Schema
t
  e :: Schema
e@Enum{TypeName
name :: TypeName
name :: Schema -> TypeName
name}           -> [(TypeName
name, Schema
e)]
  f :: Schema
f@Fixed{TypeName
name :: TypeName
name :: Schema -> TypeName
name}          -> [(TypeName
name, Schema
f)]
  Schema
_                      -> []

-- Ensures normalisation: "extracted" record is self-contained and
-- all the named types are resolvable within the scope of the schema.
normSchema :: Schema -> State (M.Map TypeName Schema) Schema
normSchema :: Schema -> State (Map TypeName Schema) Schema
normSchema Schema
r = case Schema
r of
  t :: Schema
t@(NamedType TypeName
tn) -> do
    Map TypeName Schema
resolved <- StateT (Map TypeName Schema) Identity (Map TypeName Schema)
forall s (m :: * -> *). MonadState s m => m s
get
    case TypeName -> Map TypeName Schema -> Maybe Schema
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeName
tn Map TypeName Schema
resolved of
      Just Schema
rs ->
        -- use the looked up schema (which might be a full record) and replace
        -- it in the state with NamedType for future resolves
        -- because only one full definition per schema is needed
        (Map TypeName Schema -> Map TypeName Schema)
-> StateT (Map TypeName Schema) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' (TypeName -> Schema -> Map TypeName Schema -> Map TypeName Schema
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert TypeName
tn Schema
t) StateT (Map TypeName Schema) Identity ()
-> State (Map TypeName Schema) Schema
-> State (Map TypeName Schema) Schema
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> case Schema
rs of
            NamedType TypeName
_ -> Schema -> State (Map TypeName Schema) Schema
forall (f :: * -> *) a. Applicative f => a -> f a
pure Schema
rs -- If we get a reference, the schema was already normalised.
            Schema
_ -> Schema -> State (Map TypeName Schema) Schema
normSchema Schema
rs -- Otherwise, normalise the schema before inlining.

        -- NamedType but no corresponding record?! Baaad!
      Maybe Schema
Nothing ->
        [Char] -> State (Map TypeName Schema) Schema
forall a. HasCallStack => [Char] -> a
error ([Char] -> State (Map TypeName Schema) Schema)
-> [Char] -> State (Map TypeName Schema) Schema
forall a b. (a -> b) -> a -> b
$ [Char]
"Unable to resolve schema: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Text -> [Char]
forall a. Show a => a -> [Char]
show (Schema -> Text
typeName Schema
t)

  Array Schema
s -> Schema -> Schema
Array (Schema -> Schema)
-> State (Map TypeName Schema) Schema
-> State (Map TypeName Schema) Schema
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Schema -> State (Map TypeName Schema) Schema
normSchema Schema
s
  Map Schema
s   -> Schema -> Schema
Map (Schema -> Schema)
-> State (Map TypeName Schema) Schema
-> State (Map TypeName Schema) Schema
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Schema -> State (Map TypeName Schema) Schema
normSchema Schema
s
  Union Vector Schema
l -> Vector Schema -> Schema
Union (Vector Schema -> Schema)
-> StateT (Map TypeName Schema) Identity (Vector Schema)
-> State (Map TypeName Schema) Schema
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Schema -> State (Map TypeName Schema) Schema)
-> Vector Schema
-> StateT (Map TypeName Schema) Identity (Vector Schema)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Schema -> State (Map TypeName Schema) Schema
normSchema Vector Schema
l
  r :: Schema
r@Record{name :: Schema -> TypeName
name = TypeName
tn}  -> do
    (Map TypeName Schema -> Map TypeName Schema)
-> StateT (Map TypeName Schema) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' (TypeName -> Schema -> Map TypeName Schema -> Map TypeName Schema
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert TypeName
tn (TypeName -> Schema
NamedType TypeName
tn))
    [Field]
flds <- (Field -> StateT (Map TypeName Schema) Identity Field)
-> [Field] -> StateT (Map TypeName Schema) Identity [Field]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Field
fld -> Field -> Schema -> Field
setType Field
fld (Schema -> Field)
-> State (Map TypeName Schema) Schema
-> StateT (Map TypeName Schema) Identity Field
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Schema -> State (Map TypeName Schema) Schema
normSchema (Field -> Schema
fldType Field
fld)) (Schema -> [Field]
fields Schema
r)
    Schema -> State (Map TypeName Schema) Schema
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Schema -> State (Map TypeName Schema) Schema)
-> Schema -> State (Map TypeName Schema) Schema
forall a b. (a -> b) -> a -> b
$ Schema
r { fields :: [Field]
fields = [Field]
flds }
  r :: Schema
r@Fixed{name :: Schema -> TypeName
name = TypeName
tn} -> do
    (Map TypeName Schema -> Map TypeName Schema)
-> StateT (Map TypeName Schema) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' (TypeName -> Schema -> Map TypeName Schema -> Map TypeName Schema
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert TypeName
tn (TypeName -> Schema
NamedType TypeName
tn))
    Schema -> State (Map TypeName Schema) Schema
forall (f :: * -> *) a. Applicative f => a -> f a
pure Schema
r
  r :: Schema
r@Enum{name :: Schema -> TypeName
name = TypeName
tn} -> do
    (Map TypeName Schema -> Map TypeName Schema)
-> StateT (Map TypeName Schema) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' (TypeName -> Schema -> Map TypeName Schema -> Map TypeName Schema
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert TypeName
tn (TypeName -> Schema
NamedType TypeName
tn))
    Schema -> State (Map TypeName Schema) Schema
forall (f :: * -> *) a. Applicative f => a -> f a
pure Schema
r
  Schema
s         -> Schema -> State (Map TypeName Schema) Schema
forall (f :: * -> *) a. Applicative f => a -> f a
pure Schema
s
  where
    setType :: Field -> Schema -> Field
setType Field
fld Schema
t = Field
fld { fldType :: Schema
fldType = Schema
t}