Store a file's lines in an array

Accessing a single line in `getLine_` using `List.safeIndex` takes O(n)
time with n being the number of lines. Overall this results in O(n^2)
time for formatting a diagnostic. By using an array we cut the overall
time down to O(n).

src/Error/Diagnose/Diagnostic/Internal.hs

@@ -21,15 +21,15 @@ import Data.Aeson (ToJSON(..), encode, object, (.=))
 import Data.ByteString.Lazy (ByteString)
 #endif
 
+import Data.Array (listArray)
 import Data.DList (DList)
 import qualified Data.DList as DL
 import Data.Default (Default, def)
 import Data.Foldable (fold, toList)
-import Data.HashMap.Lazy (HashMap)
 import qualified Data.HashMap.Lazy as HashMap
 import Data.List (intersperse)
 import Error.Diagnose.Report (Report)
-import Error.Diagnose.Report.Internal (prettyReport)
+import Error.Diagnose.Report.Internal (FileMap, prettyReport)
 import Error.Diagnose.Style (Annotation, Style)
 import Prettyprinter (Doc, Pretty, hardline, unAnnotate)
 import Prettyprinter.Render.Terminal (hPutDoc)
@@ -45,7 +45,7 @@ data Diagnostic msg
       -- ^ All the reports contained in a diagnostic.
       --
       --   Reports are output one by one, without connections in between.
-      (HashMap FilePath [String])
+      !FileMap
       -- ^ A map associating files with their content as lists of lines.
 
 instance Default (Diagnostic msg) where
@@ -57,7 +57,7 @@ instance Semigroup (Diagnostic msg) where
 #ifdef USE_AESON
 instance ToJSON msg => ToJSON (Diagnostic msg) where
   toJSON (Diagnostic reports files) =
-    object [ "files" .= fmap toJSONFile (HashMap.toList files)
+    object [ "files" .= fmap toJSONFile (fmap toList <$> (HashMap.toList files))
            , "reports" .= reports
            ]
     where
@@ -120,7 +120,10 @@ addFile ::
   String ->
   Diagnostic msg
 addFile (Diagnostic reports files) path content =
-  Diagnostic reports (HashMap.insert path (lines content) files)
+  let fileLines = lines content
+      lineCount = length fileLines
+      lineArray = listArray (0, lineCount - 1) fileLines
+   in Diagnostic reports (HashMap.insert path lineArray files)
 {-# INLINE addFile #-}
 
 -- | Inserts a new report into a diagnostic.

src/Error/Diagnose/Report/Internal.hs

@@ -25,9 +25,8 @@ module Error.Diagnose.Report.Internal where
 import Data.Aeson (ToJSON(..), object, (.=))
 #endif
 import Control.Applicative ((<|>))
-import Data.Array.IArray ((!))
 import qualified Data.Array.IArray as Array
-import Data.Array.Unboxed (IArray, Ix, UArray, listArray)
+import Data.Array.Unboxed (Array, IArray, Ix, UArray, listArray, (!))
 import Data.Bifunctor (bimap, first, second)
 import Data.Char.WCWidth (wcwidth)
 import Data.Default (def)
@@ -47,6 +46,8 @@ import Error.Diagnose.Style (Annotation (..))
 import Prettyprinter (Doc, Pretty (..), align, annotate, colon, hardline, lbracket, rbracket, space, width, (<+>))
 import Prettyprinter.Internal (Doc (..))
 
+type FileMap = HashMap FilePath (Array Int String)
+
 type WidthTable = UArray Int Int
 
 -- | The type of diagnostic reports with abstract message type.
@@ -158,7 +159,7 @@ err = Report True
 prettyReport ::
   Pretty msg =>
   -- | The content of the file the reports are for
-  HashMap FilePath [String] ->
+  FileMap ->
   -- | Should we print paths in unicode?
   Bool ->
   -- | The number of spaces each TAB character will span
@@ -315,7 +316,7 @@ groupMarkersPerFile markers =
 prettySubReport ::
   Pretty msg =>
   -- | The content of files in the diagnostics
-  HashMap FilePath [String] ->
+  FileMap ->
   -- | Is the output done with Unicode characters?
   Bool ->
   -- | Is the current report an error report?
@@ -372,7 +373,7 @@ splitMarkersPerLine (m@(Position {..}, _) : ms) =
 -- |
 prettyAllLines ::
   Pretty msg =>
-  HashMap FilePath [String] ->
+  FileMap ->
   Bool ->
   Bool ->
   -- | The number of spaces each TAB character will span
@@ -473,36 +474,37 @@ prettyAllLines files withUnicode isError tabSize leftLen inline multiline lineNu
 
 -- |
 getLine_ ::
-  HashMap FilePath [String] ->
+  FileMap ->
   [(Position, Marker msg)] ->
   Int ->
   Int ->
   Bool ->
   (WidthTable, Doc Annotation)
-getLine_ files markers line tabSize isError = case List.safeIndex (line - 1) =<< (HashMap.!?) files . file . fst =<< List.safeHead markers of
-  Nothing ->
-    ( mkWidthTable "",
-      annotate NoLineColor "<no line>"
-    )
-  Just code ->
-    ( mkWidthTable code,
-      flip foldMap (zip [1 ..] code) \(n, c) ->
-        let cdoc =
-              ifTab (pretty (replicate tabSize ' ')) pretty c
-            colorizingMarkers = flip filter markers \case
-              (Position (bl, bc) (el, ec) _, _)
-                | bl == el ->
-                  n >= bc && n < ec
-                | otherwise ->
-                  (bl == line && n >= bc)
-                    || (el == line && n < ec)
-                    || (bl < line && el > line)
-         in maybe
-              id
-              ((\m -> annotate (MarkerStyle $ markerColor isError m)) . snd)
-              (List.safeHead colorizingMarkers)
-              cdoc
-    )
+getLine_ files markers line tabSize isError =
+  case safeArrayIndex (line - 1) =<< (HashMap.!?) files . file . fst =<< List.safeHead markers of
+    Nothing ->
+      ( mkWidthTable "",
+        annotate NoLineColor "<no line>"
+      )
+    Just code ->
+      ( mkWidthTable code,
+        flip foldMap (zip [1 ..] code) \(n, c) ->
+          let cdoc =
+                ifTab (pretty (replicate tabSize ' ')) pretty c
+              colorizingMarkers = flip filter markers \case
+                (Position (bl, bc) (el, ec) _, _)
+                  | bl == el ->
+                    n >= bc && n < ec
+                  | otherwise ->
+                    (bl == line && n >= bc)
+                      || (el == line && n < ec)
+                      || (bl < line && el > line)
+           in maybe
+                id
+                ((\m -> annotate (MarkerStyle $ markerColor isError m)) . snd)
+                (List.safeHead colorizingMarkers)
+                cdoc
+      )
   where
     ifTab :: a -> (Char -> a) -> Char -> a
     ifTab a _ '\t' = a