Use packed representation for `Location`

3 files changed, 138 insertions, 21 deletions

diff --git a/source/Api.hs b/source/Api.hs
index 12744ce..6cffb61 100644
--- a/source/Api.hs
+++ b/source/Api.hs
@@ -111,7 +111,8 @@ findAndReadFile path = do
 lexFile :: MonadIO io => FilePath -> io (Text, [[Located Token]])
 lexFile file = do
     raw <- findAndReadFile file
-    case runLexer file raw of
+    fileId <- registerFilePath file
+    case runLexer fileId file raw of
         Left tokenError ->
             throwIO (TokenError (errorBundlePretty tokenError))
         Right (_imports, chunks) ->
diff --git a/source/Report/Location.hs b/source/Report/Location.hs
index 59e6e83..5afb319 100644
--- a/source/Report/Location.hs
+++ b/source/Report/Location.hs
@@ -1,25 +1,139 @@
 {-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE DerivingStrategies #-}
 
 module Report.Location where
 
 import Base
-import Text.Megaparsec.Pos (SourcePos (sourceColumn), sourceName, sourceLine, unPos)
+import Text.Megaparsec.Pos (SourcePos (sourceColumn, sourceLine), unPos)
+import Data.Bits
+import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef)
+import Data.IntMap.Strict qualified as IntMap
+import Data.Map.Strict qualified as Map
 import Data.Text qualified as Text
+import Data.Word (Word16)
+import System.IO.Unsafe (unsafePerformIO)
 
-data Location = Location
-    { locFile :: !FilePath
-    , locLine :: {-# UNPACK #-} !Int
-    , locColumn :: {-# UNPACK #-} !Int
-    } deriving (Show, Eq, Ord, Generic, Hashable)
+-- | File identifier used in packed source locations.
+newtype FileId = FileId
+    { unFileId :: Word16
+    } deriving stock (Show, Eq, Ord, Generic)
+      deriving anyclass (Hashable)
 
+-- | Packed source location.
+-- Bit layout (high to low):
+-- 16 bits file id | 24 bits line | 24 bits column.
+newtype Location = Location
+    { unLocation :: Word64
+    } deriving stock (Eq, Ord, Generic)
+      deriving anyclass (Hashable)
 
-fromSourcePos :: SourcePos -> Location
-fromSourcePos pos = Location
-    { locFile = sourceName pos
-    , locLine = unPos (sourceLine pos)
-    , locColumn = unPos (sourceColumn pos)
+data FileRegistry = FileRegistry
+    { pathToFileId :: Map FilePath FileId
+    , fileIdToPath :: IntMap FilePath
+    , nextFileId :: {-# UNPACK #-} !Word16
     }
 
+initialFileRegistry :: FileRegistry
+initialFileRegistry = FileRegistry
+    { pathToFileId = mempty
+    , fileIdToPath = mempty
+    , nextFileId = 0
+    }
+
+fileRegistryRef :: IORef FileRegistry
+fileRegistryRef = unsafePerformIO (newIORef initialFileRegistry)
+{-# NOINLINE fileRegistryRef #-}
+
+registerFilePath :: MonadIO io => FilePath -> io FileId
+registerFilePath path = liftIO $
+    atomicModifyIORef' fileRegistryRef \registry ->
+        case Map.lookup path (pathToFileId registry) of
+            Just fileId ->
+                (registry, fileId)
+            Nothing ->
+                if nextFileId registry == maxBound
+                    then error "registerFilePath: exhausted file-id space (16 bits)"
+                    else
+                        let fileId = FileId (nextFileId registry)
+                            fileIdInt = fromIntegral (unFileId fileId)
+                            registry' = FileRegistry
+                                { pathToFileId = Map.insert path fileId (pathToFileId registry)
+                                , fileIdToPath = IntMap.insert fileIdInt path (fileIdToPath registry)
+                                , nextFileId = nextFileId registry + 1
+                                }
+                        in
+                            (registry', fileId)
+
+lookupFilePath :: FileId -> Maybe FilePath
+lookupFilePath fileId = unsafePerformIO do
+    registry <- readIORef fileRegistryRef
+    pure (IntMap.lookup (fromIntegral (unFileId fileId)) (fileIdToPath registry))
+
+fileShift, lineShift :: Int
+fileShift = 48
+lineShift = 24
+
+fileMask, coordMask :: Word64
+fileMask = 0xFFFF
+coordMask = 0xFFFFFF
+
+nowhereWord :: Word64
+nowhereWord = maxBound
+
+mkLocation :: FileId -> Int -> Int -> Location
+mkLocation fileId line column =
+    if lineWord > coordMask || columnWord > coordMask
+        then error ("mkLocation: line/column out of range (line=" <> show line <> ", column=" <> show column <> ")")
+        else
+            Location
+                (   (fromIntegral (unFileId fileId) `shiftL` fileShift)
+                 .|. (lineWord `shiftL` lineShift)
+                 .|. columnWord
+                )
+    where
+        lineWord = fromIntegral line :: Word64
+        columnWord = fromIntegral column :: Word64
+
+locFileId :: Location -> Maybe FileId
+locFileId (Location w)
+    | w == nowhereWord = Nothing
+    | otherwise = Just (FileId (fromIntegral ((w `shiftR` fileShift) .&. fileMask)))
+
+locFile :: Location -> FilePath
+locFile loc = case locFileId loc of
+    Nothing -> "<nowhere>"
+    Just fileId ->
+        fromMaybe ("<file#" <> show (unFileId fileId) <> ">") (lookupFilePath fileId)
+
+locLine :: Location -> Int
+locLine (Location w)
+    | w == nowhereWord = -1
+    | otherwise = fromIntegral ((w `shiftR` lineShift) .&. coordMask)
+
+locColumn :: Location -> Int
+locColumn (Location w)
+    | w == nowhereWord = -1
+    | otherwise = fromIntegral (w .&. coordMask)
+
+instance Show Location where
+    showsPrec p loc =
+        showParen (p > appPrec) $
+            showString "Location {locFile = "
+            . shows (locFile loc)
+            . showString ", locLine = "
+            . shows (locLine loc)
+            . showString ", locColumn = "
+            . shows (locColumn loc)
+            . showString "}"
+      where
+        appPrec = 10
+
+fromSourcePos :: FileId -> SourcePos -> Location
+fromSourcePos fileId pos = mkLocation
+    fileId
+    (unPos (sourceLine pos))
+    (unPos (sourceColumn pos))
+
 prettyLocation :: Location -> String
 prettyLocation loc =
     locFile loc <> " " <> show (locLine loc) <> ":" <> show (locColumn loc)
@@ -32,7 +146,7 @@ class Locatable a where
     locate :: a -> Location
 
 pattern Nowhere :: Location
-pattern Nowhere = Location "<nowhere>" (-1) (-1)
+pattern Nowhere = Location 0xFFFFFFFFFFFFFFFF
 
 instance Locatable Location where
     locate = id
diff --git a/source/Syntax/Token.hs b/source/Syntax/Token.hs
index 8a40200..1ec1f7f 100644
--- a/source/Syntax/Token.hs
+++ b/source/Syntax/Token.hs
@@ -40,8 +40,8 @@ import Text.Megaparsec.Char.Lexer qualified as Lexer
 import Tptp.UnsortedFirstOrder (isAsciiLetter, isAsciiAlphaNumOrUnderscore)
 
 
-runLexer :: String -> Text -> Either (ParseErrorBundle Text Void) ([FilePath], [[Located Token]])
-runLexer file raw = runParser (evalStateT document initLexerState) file raw
+runLexer :: FileId -> String -> Text -> Either (ParseErrorBundle Text Void) ([FilePath], [[Located Token]])
+runLexer fileId file raw = runParser (evalStateT document (initLexerState fileId)) file raw
 
 
 type Lexer = StateT LexerState (Parsec Void Text)
@@ -54,14 +54,15 @@ data LexerState = LexerState
     -- to text tokens. In order to switch back to math mode correctly
     -- we need to count the braces.
     , mode :: !Mode
+    , currentFileId :: !FileId
     } deriving (Show, Eq)
 
-initLexerState :: LexerState
-initLexerState = LexerState 0 TextMode
+initLexerState :: FileId -> LexerState
+initLexerState fileId = LexerState 0 TextMode fileId
 
 incrNesting, decrNesting :: LexerState -> LexerState
-incrNesting (LexerState n m) = LexerState (succ n) m
-decrNesting (LexerState n m) = LexerState (pred n) m
+incrNesting (LexerState n m fileId) = LexerState (succ n) m fileId
+decrNesting (LexerState n m fileId) = LexerState (pred n) m fileId
 
 data Mode = TextMode | MathMode deriving (Show, Eq)
 
@@ -217,7 +218,7 @@ importBlock = do
 
 -- TODO remove once we have a proper build system and incremental compilation
 gatherImports :: Text -> [FilePath]
-gatherImports raw = case runParser (evalStateT importBlock initLexerState) "TODO filename" raw of
+gatherImports raw = case runParser (evalStateT importBlock (initLexerState (FileId maxBound))) "TODO filename" raw of
     Left err -> error (errorBundlePretty err)
     Right paths -> paths
 
@@ -507,10 +508,11 @@ closing = lexeme (group <|> optional (Char.string "\\right") *> (paren <|> brace
 -- and consumes trailing whitespace.
 lexeme :: Lexer a -> Lexer (Located a)
 lexeme p = do
+    fileId <- gets currentFileId
     start <- getSourcePos
     t <- p
     w <- whitespace
-    pure (Located (fromSourcePos start) t w)
+    pure (Located (fromSourcePos fileId start) t w)
 
 space :: Lexer Whitespace
 space = Space <$ (Char.char ' ' <|> Char.char '\n' <|> Char.char '\r')