Optimizations

src-extra/transformation/JbeamEdit/Transformation.hs

@@ -4,7 +4,7 @@
 import Data.Bool (bool)
 import Data.Foldable.Extra (notNull)
 import Data.Function (on)
 import Data.List (foldl', partition)
 import Data.List.NonEmpty (NonEmpty)
 import Data.List.NonEmpty qualified as NE
 import Data.Map (Map)
@@ -317,7 +317,7 @@
               oMap
 
       (_, listsOfNodes) = mapAccumL stepType initialMeta treesOrder
-   in V.fromList (concat listsOfNodes)
+   in foldMap V.fromList listsOfNodes
 
 treesOrder :: [VertexTreeType]
 treesOrder = [LeftTree, MiddleTree, RightTree, SupportTree]

src-extra/transformation/JbeamEdit/Transformation/BeamExtraction.hs

@@ -7,6 +7,8 @@ import Data.Map (Map)
 import Data.Map qualified as M
 import Data.Maybe (catMaybes)
 import Data.Ord (Down (Down))
+import Data.Set (Set)
+import Data.Set qualified as S
 import Data.Text (Text)
 import Data.Vector (Vector)
 import Data.Vector qualified as V
@@ -20,15 +22,14 @@ beamQuery :: NP.NodePath
 beamQuery = fromList [NP.ObjectIndex 0, NP.ObjectKey "beams"]
 
 rejectUnknownName
-  :: Foldable t
-  => t Text
+  :: Set Text
   -> Maybe (Vector Text)
   -> Maybe (Vector Text)
 rejectUnknownName knownNodeNames maybeBeam =
   bool
     maybeBeam
     Nothing
-    (any (any (`notElem` knownNodeNames)) maybeBeam)
+    (any (any (`S.notMember` knownNodeNames)) maybeBeam)
 
 possiblyBeam :: Node -> Either Node (Maybe (Vector Text))
 possiblyBeam node
@@ -53,7 +54,7 @@ vertexConns
 vertexConns maxSupport topNode vsPerType =
   go <$> extractBeams topNode
   where
-    knownNodeNames = concatMap (map anVertexName) vsPerType
+    knownNodeNames = foldMap (foldr (S.insert . anVertexName) mempty) vsPerType
     go beams =
       let possiblyInnerBeam = (:) . fmap (rejectUnknownName knownNodeNames) . possiblyBeam
           (badNodes, beamPairs) =

src/JbeamEdit/Core/Node.hs

@@ -121,7 +121,7 @@ possiblyChildren n = expectArray n <|> expectObject n
 
 moreNodesThanOne :: Vector Node -> Bool
 moreNodesThanOne v
-  | len == 1 = any moreNodesThanOne . possiblyChildren $ V.head v
+  | len == 1 = any moreNodesThanOne . possiblyChildren $ V.unsafeHead v
   | len > 1 = True
   | otherwise = False
   where

src/JbeamEdit/Formatting.hs

@@ -6,6 +6,8 @@
   formatScalarNode,
   formatNodeAndWrite,
   RuleSet (..),
+  FormattingState (..),
+  emptyFormattingState,
 ) where
 
 import Data.Bool (bool)
@@ -20,6 +22,7 @@
 import Data.Text.Encoding (encodeUtf8)
 import Data.Vector (Vector)
 import Data.Vector qualified as V
+import Data.Vector.Unboxed qualified as UV
 import JbeamEdit.Core.Node (
   InternalComment (..),
   Node (..),
@@ -64,57 +67,80 @@
 singleCharIfNot :: Char -> Bool -> Text
 singleCharIfNot a b = singleCharIf a (not b)
 
+data FormattingState = FormattingState
+  { fsUsePad :: Bool
+  , fsColumnWidths :: UV.Vector Int
+  , fsFormattedCache :: Vector (Vector Text) -- cache for data rows (without header)
+  , fsHeaderCache :: Maybe (Vector Text) -- cache for header row if it existed
+  }
+
+emptyFormattingState :: FormattingState
+emptyFormattingState = FormattingState False UV.empty V.empty Nothing
+
 addDelimiters
   :: RuleSet
   -> Int
   -> NC.NodeCursor
   -> Bool
-  -> (Bool, Vector Int) -- (usePad, columnWidths)
+  -> FormattingState
   -> [Text]
   -> [Node]
   -> [Text]
 addDelimiters _ _ _ _ _ acc [] = acc
-addDelimiters rs index c complexChildren (usePad, colWidths) acc ns@(node : rest)
+addDelimiters rs index c complexChildren fs acc ns@(node : rest)
   | complexChildren && null acc =
-      addDelimiters rs index c complexChildren (usePad, colWidths) ["\n"] ns
+      addDelimiters rs index c complexChildren fs ["\n"] ns
   | isCommentNode node =
       let formattedComment =
-            formatWithCursor
-              rs
-              (usePad, colWidths)
-              c
-              (normalizeCommentNode complexChildren node)
+            formatWithCursor rs fs c (normalizeCommentNode complexChildren node)
           formatted = (newlineBeforeComment <> formattedComment <> "\n") : acc
-       in addDelimiters rs index c complexChildren (usePad, colWidths) formatted rest
+       in addDelimiters rs index c complexChildren fs formatted rest
   | otherwise =
       case extractPreviousAssocCmt rest of
         (Just comment, rest') ->
           let baseTxt = applyCrumbAndFormat node
               formatted = (padTxt baseTxt <> " " <> formatComment comment) : acc
-           in addDelimiters
-                rs
-                (index + 1)
-                c
-                complexChildren
-                (usePad, colWidths)
-                formatted
-                rest'
+           in addDelimiters rs (index + 1) c complexChildren fs formatted rest'
         (Nothing, _) ->
           let baseTxt = applyCrumbAndFormat node
               new_acc = (padTxt baseTxt <> singleCharIf ' ' space <> singleCharIf '\n' newline) : acc
-           in addDelimiters rs (index + 1) c complexChildren (usePad, colWidths) new_acc rest
+           in addDelimiters rs (index + 1) c complexChildren fs new_acc rest
   where
     newlineBeforeComment = singleCharIfNot '\n' (any isObjectKeyNode rest || ["\n"] == acc)
 
     applyCrumbAndFormat n =
-      let padded = NC.applyCrumb c (formatWithCursor rs (usePad, colWidths)) index n
+      let padded = NC.applyCrumb c (formatWithCursor rs fs) index n
           (formatted, spaces) = splitTrailing comma padded
        in formatted <> singleCharIf ',' comma <> spaces
 
+    getCachedOrFormat nodeIdx node =
+      case (index, fsHeaderCache fs) of
+        (0, Just headerVec) ->
+          if nodeIdx < V.length headerVec
+            then headerVec V.! nodeIdx
+            else formatWithCursor rs fs c node
+        (_, Just _) ->
+          let cacheIdx = index - 1
+           in if cacheIdx >= 0 && cacheIdx < V.length (fsFormattedCache fs)
+                then
+                  let row = fsFormattedCache fs V.! cacheIdx
+                   in if nodeIdx < V.length row
+                        then row V.! nodeIdx
+                        else formatWithCursor rs fs c node
+                else formatWithCursor rs fs c node
+        (_, Nothing) ->
+          if index < V.length (fsFormattedCache fs)
+            then
+              let row = fsFormattedCache fs V.! index
+               in if nodeIdx < V.length row
+                    then row V.! nodeIdx
+                    else formatWithCursor rs fs c node
+            else formatWithCursor rs fs c node
+
     padTxt baseTxt =
-      if usePad && not (isCommentNode node) && comma
+      if fsUsePad fs && not (isCommentNode node) && comma
         then
-          let width = sum (colWidths V.!? index)
+          let width = sum (fsColumnWidths fs UV.!? index)
            in T.justifyLeft (width + 1) ' ' baseTxt
         else baseTxt
 
@@ -127,53 +153,72 @@
   | T.all isSpace s = s
   | otherwise = T.replicate n " " <> s
 
-skipHeaderRow :: Vector (Vector Node) -> Vector (Vector Node)
-skipHeaderRow nodes =
+extractHeader
+  :: Vector (Vector Node) -> (Maybe (Vector Node), Vector (Vector Node))
+extractHeader nodes =
   case V.uncons nodes of
     Just (headerRow, rest) ->
-      bool nodes rest (all isStringNode headerRow)
-    Nothing -> nodes
+      if all isStringNode headerRow
+        then (Just headerRow, rest)
+        else (Nothing, nodes)
+    Nothing -> (Nothing, nodes)
 
-maxColumnLengths
-  :: RuleSet -> NC.NodeCursor -> Vector (Vector Node) -> Vector Int
-maxColumnLengths rs cursor rows
-  | V.null rows = V.empty
+maxColumnLengthsWithCache
+  :: RuleSet
+  -> NC.NodeCursor
+  -> Vector (Vector Node)
+  -> (Maybe (Vector Text), UV.Vector Int, Vector (Vector Text))
+maxColumnLengthsWithCache rs cursor rows
+  | V.null rows = (Nothing, UV.empty, V.empty)
   | otherwise =
-      V.map
-        (V.maximum . V.map T.length)
-        (transposeWithPadding rs cursor $ skipHeaderRow rows)
+      let (mHeader, dataRows) = extractHeader rows
+          headerFormatted = V.map (formatWithCursor rs emptyFormattingState cursor) <$> mHeader
+          formatted = transposeWithPadding rs cursor dataRows
+          colLengths =
+            if V.null formatted
+              then UV.empty
+              else
+                UV.fromList $
+                  V.foldr (\col acc -> V.maximum (V.map T.length col) : acc) [] formatted
+       in (headerFormatted, colLengths, formatted)
 
 transposeWithPadding
-  :: RuleSet -> NC.NodeCursor -> Vector (Vector Node) -> Vector (Vector T.Text)
+  :: RuleSet -> NC.NodeCursor -> Vector (Vector Node) -> Vector (Vector Text)
 transposeWithPadding rs cursor vvs =
-  let numCols = V.maximum (V.map V.length vvs)
-   in V.generate numCols $ \j ->
-        V.map
-          ( \row ->
-              mwhen
-                (j < V.length row)
-                (formatWithCursor rs (False, V.empty) cursor (row V.! j))
-          )
-          vvs
+  let numCols = if V.null vvs then 0 else V.maximum (V.map V.length vvs)
+   in if numCols == 0
+        then V.empty
+        else V.generate numCols $ \j ->
+          V.map
+            ( \row ->
+                mwhen
+                  (j < V.length row)
+                  (formatWithCursor rs emptyFormattingState cursor (row V.! j))
+            )
+            vvs
 
 doFormatNode
   :: RuleSet
   -> NC.NodeCursor
-  -> (Bool, Vector Int)
+  -> FormattingState
   -> Vector Node
   -> Text
-doFormatNode rs cursor padAmounts nodes =
+doFormatNode rs cursor _ nodes =
   let autoPadEnabled =
         lookupPropertyForCursor ExactMatch AutoPad rs cursor == Just True
 
       childrenVectors = V.map (fromMaybe V.empty . expectArray) nodes
-      padAmounts' = maxColumnLengths rs cursor childrenVectors
-      maybePadAmounts =
-        bool padAmounts (False, padAmounts') autoPadEnabled
+      (mHeaderFormatted, colWidths, formattedCache) =
+        maxColumnLengthsWithCache rs cursor childrenVectors
+
+      fs' =
+        if autoPadEnabled
+          then FormattingState True colWidths formattedCache mHeaderFormatted
+          else emptyFormattingState
 
       formatted =
         reverse
-          . addDelimiters rs 0 cursor complexChildren maybePadAmounts []
+          . addDelimiters rs 0 cursor complexChildren fs' []
           . V.toList
           $ nodes
 
@@ -213,44 +258,39 @@
 formatScalarNode _ = error "Unhandled scalar node"
 
 formatWithCursor
-  :: RuleSet -> (Bool, Vector Int) -> NC.NodeCursor -> Node -> Text
-formatWithCursor rs (_, maybePadAmounts) cursor (Array a)
+  :: RuleSet -> FormattingState -> NC.NodeCursor -> Node -> Text
+formatWithCursor rs fs cursor (Array a)
   | V.null a = "[]"
   | otherwise =
       T.concat
         [ "["
-        , doFormatNode rs cursor (notNull maybePadAmounts, maybePadAmounts) a
+        , doFormatNode rs cursor fs a
         , "]"
         ]
-formatWithCursor rs (_, maybePadAmounts) cursor (Object o)
+formatWithCursor rs fs cursor (Object o)
   | V.null o = "{}"
   | otherwise =
       T.concat
         [ "{"
-        , doFormatNode rs cursor (notNull maybePadAmounts, maybePadAmounts) o
+        , doFormatNode rs cursor fs o
         , "}"
         ]
-formatWithCursor rs (_, maybePadAmounts) cursor (ObjectKey (k, v)) =
+formatWithCursor rs fs cursor (ObjectKey (k, v)) =
   T.concat
-    [ formatWithCursor rs (notNull maybePadAmounts, maybePadAmounts) cursor k
+    [ formatWithCursor rs fs cursor k
     , " : "
-    , formatWithCursor rs (notNull maybePadAmounts, maybePadAmounts) cursor v
+    , formatWithCursor rs fs cursor v
     ]
 formatWithCursor _ _ _ (Comment comment) = formatComment comment
 formatWithCursor rs _ cursor n =
   let ps = findPropertiesForCursor PrefixMatch cursor rs
    in applyPadLogic formatScalarNode ps n
 
 formatNode :: RuleSet -> Node -> Text
-formatNode rs node = formatWithCursor rs (False, V.empty) newCursor node <> T.singleton '\n'
+formatNode rs node = formatWithCursor rs emptyFormattingState newCursor node <> T.singleton '\n'
 
-#ifdef ENABLE_WINDOWS_NEWLINES
-replaceNewlines :: Text -> Text
-replaceNewlines = T.replace "\n" "\r\n"
-#else
 replaceNewlines :: Text -> Text
 replaceNewlines = id
-#endif
 
 formatNodeAndWrite
   :: RuleSet

test/FormattingSpec.hs

@@ -87,7 +87,7 @@ spec = do
       applySpecOnInput
         descFun
         shouldBe
-        (formatWithCursor mempty (False, mempty) newCursor node)
+        (formatWithCursor mempty emptyFormattingState newCursor node)
         (T.pack jbeam)
     descFun jbeam node = "should format " ++ show node ++ " as " ++ jbeam
     specs =