3 files changed, 88 insertions, 32 deletions

diff --git a/src/Data/JLD.hs b/src/Data/JLD.hs
index d7688d0..a7042dc 100644
--- a/src/Data/JLD.hs
+++ b/src/Data/JLD.hs
@@ -4,6 +4,7 @@ module Data.JLD (
     module Data.JLD.Options,
     JLDExpansionParams (..),
     JLDExpansionState (..),
+    JLDFlatteningParams (..),
     expand,
     flatten,
 ) where
@@ -100,9 +101,44 @@ expand document baseUrl paramsFn = do
 
     pure (result', state')
 
-flatten :: Monad m => Value -> URI -> (JLDExpansionParams () m -> JLDExpansionParams e m) -> m (Either (JLDError e) Value, JLDExpansionState)
+data JLDFlatteningParams e m = JLDFlatteningParams
+    { jldFlatteningParamsDocumentLoader :: DocumentLoader e m
+    , jldFlatteningParamsProcessingMode :: JLDVersion
+    , jldFlatteningParamsMaxRemoteContexts :: Int
+    , jldFlatteningParamsExpandContext :: Maybe Value
+    , jldFlatteningParamsFrameExpansion :: Bool
+    , jldFlatteningParamsState :: JLDExpansionState
+    , jldFlatteningParamsCompactArrays :: Bool
+    }
+    deriving (Show)
+
+flatten :: Monad m => Value -> URI -> (JLDFlatteningParams () m -> JLDFlatteningParams e m) -> m (Either (JLDError e) Value, JLDExpansionState)
 flatten document baseUrl paramsFn = do
-    (result, state') <- expand document baseUrl paramsFn
+    let JLDFlatteningParams{..} =
+            paramsFn
+                JLDFlatteningParams
+                    { jldFlatteningParamsDocumentLoader = DocumentLoader <. const <. pure <| Left ()
+                    , jldFlatteningParamsProcessingMode = JLD1_1
+                    , jldFlatteningParamsMaxRemoteContexts = 20
+                    , jldFlatteningParamsExpandContext = Nothing
+                    , jldFlatteningParamsFrameExpansion = False
+                    , jldFlatteningParamsState =
+                        JLDExpansionState
+                            { jldExpansionStateContextCache = mempty
+                            , jldExpansionStateDocumentCache = mempty
+                            }
+                    , jldFlatteningParamsCompactArrays = True
+                    }
+        expansionParams =
+            JLDExpansionParams
+                { jldExpansionParamsDocumentLoader = jldFlatteningParamsDocumentLoader
+                , jldExpansionParamsProcessingMode = jldFlatteningParamsProcessingMode
+                , jldExpansionParamsMaxRemoteContexts = jldFlatteningParamsMaxRemoteContexts
+                , jldExpansionParamsExpandContext = jldFlatteningParamsExpandContext
+                , jldExpansionParamsFrameExpansion = jldFlatteningParamsFrameExpansion
+                , jldExpansionParamsState = jldFlatteningParamsState
+                }
+    (result, state') <- expand document baseUrl (const expansionParams)
     case result of
         Left err -> pure (Left err, state')
         Right expanded -> fmap (,state') <. runExceptT <| F.flatten expanded
diff --git a/src/Data/JLD/Flattening.hs b/src/Data/JLD/Flattening.hs
index 2bfd8dd..3a8c726 100644
--- a/src/Data/JLD/Flattening.hs
+++ b/src/Data/JLD/Flattening.hs
@@ -34,7 +34,7 @@ collectNodesStep ar node
 flatten :: Monad m => Value -> JLDFlatteningT e m Value
 flatten element = do
     -- 1. 2.
-    nodeMap <- fst <$> buildNodeMap element id
+    (nodeMap, _, _, _) <- buildNodeMap element id
 
     -- 3. 4.
     let defaultGraph = fromMaybe mempty <| M.lookup (show KeywordDefault) nodeMap
diff --git a/src/Data/JLD/Flattening/NodeMap.hs b/src/Data/JLD/Flattening/NodeMap.hs
index 65db9ab..6c35302 100644
--- a/src/Data/JLD/Flattening/NodeMap.hs
+++ b/src/Data/JLD/Flattening/NodeMap.hs
@@ -19,6 +19,7 @@ import Data.Aeson.KeyMap qualified as KM (filterWithKey, insert, lookup, member,
 import Data.Foldable.WithIndex (FoldableWithIndex (..), iforM_)
 import Data.Map.Strict qualified as M (insert, lookup)
 import Data.Vector qualified as V (singleton, snoc, uniq)
+import Debug.Pretty.Simple (pTraceShowM)
 
 type BNMT e m = REST BNMEnv (Either (JLDError e) ()) BNMState m
 
@@ -45,6 +46,8 @@ data BNMParams = BNMParams
     , bnmParamsActiveProperty :: Maybe Text
     , bnmParamsList :: Maybe Array
     , bnmParamsReferenceNode :: Maybe Object
+    , bnmParamsIdentifierCounter :: Int
+    , bnmParamsIdentifierMap :: Map Text Text
     }
     deriving (Show, Eq)
 
@@ -78,19 +81,26 @@ bnmBuildNodeMap :: Monad m => Value -> (BNMParams -> BNMParams) -> BNMT e m (May
 bnmBuildNodeMap value paramsFn = do
     BNMEnv{..} <- ask
     nodeMap <- gets bnmStateNodeMap
+    identifierCounter <- gets bnmStateIdentifierCounter
+    identifierMap <- gets bnmStateIdentifierMap
     let params p =
             paramsFn
                 p
                     { bnmParamsNodeMap = nodeMap
                     , bnmParamsActiveGraph = bnmEnvActiveGraph
-                    , bnmParamsActiveSubject = bnmEnvActiveSubject
-                    , bnmParamsActiveProperty = bnmEnvActiveProperty
+                    , bnmParamsIdentifierCounter = identifierCounter
+                    , bnmParamsIdentifierMap = identifierMap
                     }
-    (nodeMap', list) <-
+    (nodeMap', list, identifierCounter', identifierMap') <-
         buildNodeMap value params |> runExceptT >=> \case
             Left err -> throwError <| Left err
             Right a -> pure a
-    bnmModifyNodeMap <| const nodeMap'
+    modify \st ->
+        st
+            { bnmStateNodeMap = nodeMap'
+            , bnmStateIdentifierCounter = identifierCounter'
+            , bnmStateIdentifierMap = identifierMap'
+            }
     pure list
 
 buildNodeMap' :: Monad m => Value -> BNMT e m ()
@@ -138,7 +148,14 @@ buildNodeMap' element = case element of
                 -- 5.
                 | Just elemList <- KM.lookup (show KeywordList) elementObject' -> do
                     -- 5.1. 5.2.
-                    subList <- listToObject <$> bnmBuildNodeMap elemList id
+                    subList <-
+                        listToObject <$> bnmBuildNodeMap elemList \params ->
+                            params
+                                { bnmParamsActiveSubject = bnmEnvActiveSubject
+                                , bnmParamsActiveProperty = bnmEnvActiveProperty
+                                , bnmParamsList = Just mempty
+                                , bnmParamsReferenceNode = bnmEnvReferenceNode
+                                }
 
                     gets bnmStateList >>= \case
                         -- 5.3.
@@ -152,24 +169,25 @@ buildNodeMap' element = case element of
                 | isNodeObject (Object elementObject') -> do
                     id' <- case KM.lookup (show KeywordId) elementObject' of
                         -- 6.1.
-                        Just (String idValue) | isBlankIri idValue -> bnmCreateIdentifier <| Just idValue
+                        Just (String idValue)
+                            | isBlankIri idValue -> bnmCreateIdentifier (Just idValue)
+                            | otherwise -> pure idValue
                         Just _ -> throwError <| Right ()
                         -- 6.2.
                         Nothing -> bnmCreateIdentifier Nothing
 
                     -- 6.3.
-                    nodeMap <- gets bnmStateNodeMap
-                    when (N.hasKey2 bnmEnvActiveGraph (Just id') nodeMap) do
-                        bnmModifyNodeMap <| N.insert bnmEnvActiveGraph (Just id') (Just <| show KeywordId) (String id')
+                    bnmModifyNodeMap \nodeMap ->
+                        if N.hasKey2 bnmEnvActiveGraph (Just id') nodeMap
+                            then nodeMap
+                            else N.insert bnmEnvActiveGraph (Just id') (Just <| show KeywordId) (String id') nodeMap
 
-                    nodeMap' <- gets bnmStateNodeMap
                     if
                             -- 6.5.
-                            | Just referenceNode <- bnmEnvReferenceNode ->
-                                unless (N.memberArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty (Object referenceNode) nodeMap') do
-                                    bnmModifyNodeMap
-                                        <. N.modifyArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty
-                                        <| flip V.snoc (Object referenceNode)
+                            | Just referenceNode <- bnmEnvReferenceNode -> bnmModifyNodeMap \nodeMap ->
+                                if N.memberArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty (Object referenceNode) nodeMap
+                                    then nodeMap
+                                    else N.modifyArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty (`V.snoc` Object referenceNode) nodeMap
                             -- 6.6.
                             | isJust bnmEnvActiveProperty -> do
                                 -- 6.6.1.
@@ -177,11 +195,10 @@ buildNodeMap' element = case element of
 
                                 gets bnmStateList >>= \case
                                     -- 6.6.2.
-                                    Nothing ->
-                                        unless (N.memberArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty reference nodeMap') do
-                                            bnmModifyNodeMap
-                                                <. N.modifyArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty
-                                                <| flip V.snoc reference
+                                    Nothing -> bnmModifyNodeMap \nodeMap ->
+                                        if N.memberArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty reference nodeMap
+                                            then nodeMap
+                                            else N.modifyArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty (`V.snoc` reference) nodeMap
                                     -- 6.6.3.
                                     Just stateList -> bnmModifyList <. const <. Just <| V.snoc stateList reference
                             --
@@ -202,10 +219,10 @@ buildNodeMap' element = case element of
                         Nothing -> pure ()
 
                     -- 6.8.
-                    nodeMap'' <- gets bnmStateNodeMap
+                    nodeMap' <- gets bnmStateNodeMap
                     case KM.lookup (show KeywordIndex) elementObject' of
                         Just indexValue
-                            | N.hasKey3 bnmEnvActiveGraph (Just id') (Just <| show KeywordIndex) nodeMap'' -> throwError <| Left ConflictingIndexes
+                            | N.hasKey3 bnmEnvActiveGraph (Just id') (Just <| show KeywordIndex) nodeMap' -> throwError <| Left ConflictingIndexes
                             | otherwise -> bnmModifyNodeMap <| N.insert bnmEnvActiveGraph (Just id') (Just <| show KeywordType) indexValue
                         --
                         Nothing -> pure ()
@@ -257,9 +274,10 @@ buildNodeMap' element = case element of
                                         else pure property
 
                                 -- 6.12.2.
-                                gets (bnmStateNodeMap .> N.hasKey3 bnmEnvActiveGraph (Just id') (Just property')) >>= \case
-                                    True -> pure ()
-                                    False -> bnmModifyNodeMap <| N.insert bnmEnvActiveGraph (Just id') (Just property') (Array mempty)
+                                bnmModifyNodeMap \nodeMap ->
+                                    if N.hasKey3 bnmEnvActiveGraph (Just id') (Just property') nodeMap
+                                        then nodeMap
+                                        else N.insert bnmEnvActiveGraph (Just id') (Just property') (Array mempty) nodeMap
 
                                 void <| bnmBuildNodeMap value \params ->
                                     params
@@ -273,12 +291,12 @@ buildNodeMap' element = case element of
     --
     _ -> pure ()
 
-buildNodeMap :: Monad m => Value -> (BNMParams -> BNMParams) -> JLDFlatteningT e m (NodeMap, Maybe Array)
+buildNodeMap :: Monad m => Value -> (BNMParams -> BNMParams) -> JLDFlatteningT e m (NodeMap, Maybe Array, Int, Map Text Text)
 buildNodeMap document paramsFn = do
     (result, BNMState{..}) <- buildNodeMap' document |> runREST env st
     case result of
         Left (Left err) -> throwError err
-        _ -> pure (bnmStateNodeMap, bnmStateList)
+        _ -> pure (bnmStateNodeMap, bnmStateList, bnmStateIdentifierCounter, bnmStateIdentifierMap)
   where
     BNMParams{..} =
         paramsFn
@@ -289,6 +307,8 @@ buildNodeMap document paramsFn = do
                 , bnmParamsActiveProperty = Nothing
                 , bnmParamsList = mempty
                 , bnmParamsReferenceNode = Nothing
+                , bnmParamsIdentifierCounter = 0
+                , bnmParamsIdentifierMap = mempty
                 }
 
     env =
@@ -303,8 +323,8 @@ buildNodeMap document paramsFn = do
         BNMState
             { bnmStateNodeMap = bnmParamsNodeMap
             , bnmStateList = bnmParamsList
-            , bnmStateIdentifierCounter = 1
-            , bnmStateIdentifierMap = mempty
+            , bnmStateIdentifierCounter = bnmParamsIdentifierCounter
+            , bnmStateIdentifierMap = bnmParamsIdentifierMap
             }
 
 mergeNodeMaps :: NodeMap -> NodeMap