8 files changed, 153 insertions, 70 deletions

diff --git a/jsonld.cabal b/jsonld.cabal
index 0efa3b2..376a647 100644
--- a/jsonld.cabal
+++ b/jsonld.cabal
@@ -75,6 +75,7 @@ library
     , megaparsec
     , modern-uri
     , mtl
+    , pretty-simple
     , rdf4h
     , relude
     , req
@@ -120,6 +121,7 @@ test-suite jsonld-test
     , megaparsec
     , modern-uri
     , mtl
+    , pretty-simple
     , rdf4h
     , relude
     , req
diff --git a/package.yaml b/package.yaml
index 60226fb..adc0acb 100644
--- a/package.yaml
+++ b/package.yaml
@@ -28,6 +28,7 @@ dependencies:
     - megaparsec
     - modern-uri
     - mtl
+    - pretty-simple
     - rdf4h
     - relude
     - req
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
diff --git a/test/Test/Common.hs b/test/Test/Common.hs
index ffc3264..e386551 100644
--- a/test/Test/Common.hs
+++ b/test/Test/Common.hs
@@ -1,24 +1,20 @@
-module Test.Common (W3CTestList (..), W3CTest (..), W3CTestOption (..), fetchTest, parseExpansionOptions) where
+module Test.Common (W3CTestList (..), W3CTest (..), W3CTestOption (..), documentLoader, fetchTest) where
 
 import Data.JLD.Prelude
 
-import Test.Tasty
-import Test.Tasty.ExpectedFailure (ignoreTestBecause)
-import Test.Tasty.HUnit (assertFailure, testCase, (@?=))
-
 import Data.Aeson (FromJSON (..), Value (..), (.:), (.:?))
 import Data.Aeson.Types (prependFailure, typeMismatch)
-import Data.JLD (DocumentLoader (..), JLDExpansionParams (..), JLDVersion (..), mimeType, toJldErrorCode)
-import Data.JLD.Model.URI (parseUri)
+import Data.JLD (DocumentLoader (..), mimeType)
 import Data.Maybe (fromJust)
 import Network.HTTP.Req (GET (..), NoReqBody (..), defaultHttpConfig, header, jsonResponse, req, responseBody, runReq, useHttpsURI, useURI)
-import Text.URI (URI, mkURI, relativeTo)
+import Text.URI (URI)
 
 data W3CTestOption = W3CTestOption
     { w3cTestOptionSpecVersion :: Maybe Text
     , w3cTestOptionProcessingMode :: Maybe Text
     , w3cTestOptionBase :: Maybe Text
     , w3cTestOptionExpandContext :: Maybe Text
+    , w3cTestOptionCompactArrays :: Maybe Bool
     }
     deriving (Show)
 
@@ -29,6 +25,7 @@ instance FromJSON W3CTestOption where
             <*> (v .:? "processingMode")
             <*> (v .:? "base")
             <*> (v .:? "expandContext")
+            <*> (v .:? "compactArrays")
     parseJSON invalid = prependFailure "parsing W3CTestOption failed, " (typeMismatch "Object" invalid)
 
 data W3CTest = W3CTest
@@ -72,25 +69,3 @@ fetchTest url = do
     runReq defaultHttpConfig do
         res <- req GET reqUrl NoReqBody jsonResponse (reqOptions <> header "Accept" mimeType)
         pure <| responseBody res
-
-parseExpansionOptions :: URI -> URI -> Maybe W3CTestOption -> IO (URI, JLDExpansionParams () IO -> JLDExpansionParams Text IO)
-parseExpansionOptions baseUrl inputUrl maybeOptions = do
-    expandContext <- case maybeOptions >>= w3cTestOptionExpandContext of
-        Just rawUrl -> do
-            url <- fromJust <. (`relativeTo` baseUrl) <$> mkURI rawUrl
-            Just <$> fetchTest url
-        Nothing -> pure Nothing
-
-    let params p =
-            p
-                { jldExpansionParamsDocumentLoader = documentLoader
-                , jldExpansionParamsProcessingMode = case maybeOptions >>= w3cTestOptionProcessingMode of
-                    Just "json-ld-1.0" -> JLD1_0
-                    Just "json-ld-1.1" -> JLD1_1
-                    _ -> jldExpansionParamsProcessingMode p
-                , jldExpansionParamsExpandContext = expandContext <|> jldExpansionParamsExpandContext p
-                }
-
-    pure (expandBaseUrl, params)
-  where
-    expandBaseUrl = fromMaybe inputUrl (parseUri =<< w3cTestOptionBase =<< maybeOptions)
diff --git a/test/Test/Expansion.hs b/test/Test/Expansion.hs
index b5b1e07..0d553a7 100644
--- a/test/Test/Expansion.hs
+++ b/test/Test/Expansion.hs
@@ -6,10 +6,33 @@ import Test.Tasty
 import Test.Tasty.ExpectedFailure (ignoreTestBecause)
 import Test.Tasty.HUnit
 
-import Data.JLD (expand, toJldErrorCode)
+import Data.JLD (JLDExpansionParams (..), JLDVersion (..), expand, toJldErrorCode)
+import Data.JLD.Model.URI (parseUri)
 import Data.Maybe (fromJust)
-import Test.Common (W3CTest (..), W3CTestList (..), W3CTestOption (..), fetchTest, parseExpansionOptions)
-import Text.URI (mkURI, relativeTo)
+import Test.Common (W3CTest (..), W3CTestList (..), W3CTestOption (..), documentLoader, fetchTest)
+import Text.URI (URI, mkURI, relativeTo)
+
+parseExpansionOptions :: URI -> URI -> Maybe W3CTestOption -> IO (URI, JLDExpansionParams () IO -> JLDExpansionParams Text IO)
+parseExpansionOptions baseUrl inputUrl maybeOptions = do
+    expandContext <- case maybeOptions >>= w3cTestOptionExpandContext of
+        Just rawUrl -> do
+            url <- fromJust <. (`relativeTo` baseUrl) <$> mkURI rawUrl
+            Just <$> fetchTest url
+        Nothing -> pure Nothing
+
+    let params p =
+            p
+                { jldExpansionParamsDocumentLoader = documentLoader
+                , jldExpansionParamsProcessingMode = case maybeOptions >>= w3cTestOptionProcessingMode of
+                    Just "json-ld-1.0" -> JLD1_0
+                    Just "json-ld-1.1" -> JLD1_1
+                    _ -> jldExpansionParamsProcessingMode p
+                , jldExpansionParamsExpandContext = expandContext <|> jldExpansionParamsExpandContext p
+                }
+
+    pure (expandBaseUrl, params)
+  where
+    expandBaseUrl = fromMaybe inputUrl (parseUri =<< w3cTestOptionBase =<< maybeOptions)
 
 expansionTests :: W3CTestList -> TestTree
 expansionTests testList = testGroup "Expansion" <| uncurry expansionTest <$> zip (w3cSequence testList) [1 ..]
diff --git a/test/Test/Flattening.hs b/test/Test/Flattening.hs
index bc64b88..76f5434 100644
--- a/test/Test/Flattening.hs
+++ b/test/Test/Flattening.hs
@@ -6,10 +6,36 @@ import Test.Tasty
 import Test.Tasty.ExpectedFailure (ignoreTestBecause)
 import Test.Tasty.HUnit
 
-import Data.JLD (expand, flatten, toJldErrorCode)
+import Data.JLD (JLDFlatteningParams (..), JLDVersion (..), flatten, toJldErrorCode)
+import Data.JLD.Model.URI (parseUri)
 import Data.Maybe (fromJust)
-import Test.Common (W3CTest (..), W3CTestList (..), W3CTestOption (..), fetchTest, parseExpansionOptions)
-import Text.URI (mkURI, relativeTo)
+import Test.Common (W3CTest (..), W3CTestList (..), W3CTestOption (..), documentLoader, fetchTest)
+import Text.URI (URI, mkURI, relativeTo)
+
+parseFlatteningOptions :: URI -> URI -> Maybe W3CTestOption -> IO (URI, JLDFlatteningParams () IO -> JLDFlatteningParams Text IO)
+parseFlatteningOptions baseUrl inputUrl maybeOptions = do
+    expandContext <- case maybeOptions >>= w3cTestOptionExpandContext of
+        Just rawUrl -> do
+            url <- fromJust <. (`relativeTo` baseUrl) <$> mkURI rawUrl
+            Just <$> fetchTest url
+        Nothing -> pure Nothing
+
+    let params p =
+            p
+                { jldFlatteningParamsDocumentLoader = documentLoader
+                , jldFlatteningParamsProcessingMode = case maybeOptions >>= w3cTestOptionProcessingMode of
+                    Just "json-ld-1.0" -> JLD1_0
+                    Just "json-ld-1.1" -> JLD1_1
+                    _ -> jldFlatteningParamsProcessingMode p
+                , jldFlatteningParamsExpandContext = expandContext <|> jldFlatteningParamsExpandContext p
+                , jldFlatteningParamsCompactArrays = case maybeOptions >>= w3cTestOptionCompactArrays of
+                    Just b -> b
+                    _ -> jldFlatteningParamsCompactArrays p
+                }
+
+    pure (expandBaseUrl, params)
+  where
+    expandBaseUrl = fromMaybe inputUrl (parseUri =<< w3cTestOptionBase =<< maybeOptions)
 
 flatteningTests :: W3CTestList -> TestTree
 flatteningTests testList = testGroup "Flattening" <| uncurry flatteningTest <$> zip (w3cSequence testList) [1 ..]
@@ -29,7 +55,7 @@ flatteningTest W3CTest{..} (show .> (<> ". " <> toString w3cTestName) -> testNam
             inputJld <- fetchTest inputUrl
             expectJld <- fetchTest expectUrl
 
-            (expandBaseUrl, params) <- parseExpansionOptions baseUrl inputUrl w3cTestOption
+            (expandBaseUrl, params) <- parseFlatteningOptions baseUrl inputUrl w3cTestOption
             (result, _) <- flatten inputJld expandBaseUrl params
 
             result @?= Right expectJld
@@ -41,7 +67,7 @@ flatteningTest W3CTest{..} (show .> (<> ". " <> toString w3cTestName) -> testNam
 
             inputJld <- fetchTest inputUrl
 
-            (expandBaseUrl, params) <- parseExpansionOptions baseUrl inputUrl w3cTestOption
+            (expandBaseUrl, params) <- parseFlatteningOptions baseUrl inputUrl w3cTestOption
             (result, _) <- flatten inputJld expandBaseUrl params
 
             (result |> first toJldErrorCode) @?= Left expectErrorRaw