Implement Box instance from Layer to Object, either all at once or indexed by an ObjectId

src/PDF/Layer.hs

@@ -1,19 +1,28 @@
 {-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
 module PDF.Layer (
       Layer(..)
     , unify
   ) where
 
+import Control.Monad.Fail (MonadFail(..))
 import Data.Map (Map, (!), mapWithKey, member)
-import qualified Data.Map as Map (empty, union)
+import qualified Data.Map as Map (
+    empty, insert, keysSet, lookup, member, union
+  )
+import qualified Data.Set as Set (delete, toList)
+import PDF.Box (Box(..))
 import PDF.Object (
-      IndexedObjects, IndirectObjCoordinates(..), Occurrence(..)
+      IndexedObjects, IndirectObjCoordinates(..), Object, Occurrence(..)
     , Structure(..), XRefEntry(..), XRefSection, eofMarker, outputBody
   )
 import qualified PDF.Output as Output (line)
 import PDF.Output (
-    Offset, Output(..), Resource(..), byteString, getOffset, getOffsets, newLine
+      ObjectId, Offset(..), Output(..), Resource(..), byteString, getOffset
+    , getOffsets, newLine
   )
+import Prelude hiding (fail)
 import Text.Printf (printf)
 
 data Layer = Layer {
@@ -44,6 +53,38 @@ instance Output Layer where
     where
       Structure {xRef, trailer} = docStructure
 
+data AllObjects = AllObjects
+
+instance MonadFail m => Box m AllObjects Layer (Map ObjectId Object) where
+  r AllObjects = return . objects
+  w AllObjects layer@(Layer {occurrences, docStructure}) newObjects =
+    return $ layer {
+          occurrences = keptOccurrences ++ newOccurrences
+        , objects = newObjects
+        , docStructure = docStructure {
+              xRef = (const $ InUse (Offset 0) 0) <$> newObjects
+            }
+      }
+    where
+      filterOccurrences c@(Comment _) (occ, newObjIds) = (c:occ, newObjIds)
+      filterOccurrences i@(Indirect (IndirectObjCoordinates {objectId})) (occ, newObjIds)
+        | Map.member objectId newObjects = (i:occ, Set.delete objectId newObjIds)
+        | otherwise = (occ, newObjIds)
+      (keptOccurrences, newObjectIds) =
+        foldr filterOccurrences ([], Map.keysSet newObjects) occurrences
+      makeOccurrence objectId =
+        Indirect (IndirectObjCoordinates {objectId, versionNumber = 0})
+      newOccurrences = makeOccurrence <$> Set.toList newObjectIds
+
+instance MonadFail m => Box m ObjectId Layer Object where
+  r objectId = maybe (fail errorMsg) return . Map.lookup objectId . objects
+    where
+      errorMsg = "Unknown " ++ show objectId
+  w objectId layer@(Layer {objects}) object
+    | Map.member objectId objects =
+        return $ layer {objects = Map.insert objectId object objects}
+    | otherwise = fail $ "Unknown " ++ show objectId
+
 emptyLayer :: Layer
 emptyLayer = Layer {
       docStructure = Structure {xRef = Map.empty, trailer = Map.empty}