-- | Conversion between code-point and UTF-16 positions at the LSP boundary.
--
-- LSP positions count UTF-16 code units (the default position encoding, and
-- the only one VS Code supports), while alex positions and everything derived
-- from them (the surface AST, the reference index, the typechecker) count
-- Unicode code points. The two agree except on lines containing astral-plane
-- characters (code points above U+FFFF, such as @𝕀@), each of which takes two
-- UTF-16 units. Emitting a code-point column for such a line shifts every
-- position to the right of the character, and a range boundary can land in
-- the middle of a surrogate pair, which is what shatters the glyph in the
-- editor (see issue #303).
--
-- Everything internal stays in code points; the handlers convert with this
-- module when crossing the wire, in both directions. 'AstralLines' records
-- only the lines where the two encodings differ, so for the common all-BMP
-- document every conversion is an identity after one map lookup.
module Language.Rzk.VSCode.PositionEncoding (
  AstralLines,
  astralLines,
  utf16Length,
  colToUtf16,
  colFromUtf16,
  positionFromUtf16,
  rangeToUtf16,
  tokensToUtf16,
) where

import qualified Data.IntMap.Strict          as IntMap
import qualified Data.Text                   as T
import           Language.LSP.Protocol.Types (Position (Position),
                                              Range (Range),
                                              SemanticTokenAbsolute (..))

-- | The lines of a document on which code-point and UTF-16 columns differ,
-- keyed by 0-based line number.
newtype AstralLines = AstralLines (IntMap.IntMap T.Text)

astralLines :: T.Text -> AstralLines
astralLines :: Text -> AstralLines
astralLines Text
src = IntMap Text -> AstralLines
AstralLines (IntMap Text -> AstralLines) -> IntMap Text -> AstralLines
forall a b. (a -> b) -> a -> b
$ [(Int, Text)] -> IntMap Text
forall a. [(Int, a)] -> IntMap a
IntMap.fromDistinctAscList
  [ (Int
i, Text
line) | (Int
i, Text
line) <- [Int] -> [Text] -> [(Int, Text)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (Text -> [Text]
T.lines Text
src), (Char -> Bool) -> Text -> Bool
T.any Char -> Bool
isAstral Text
line ]

isAstral :: Char -> Bool
isAstral :: Char -> Bool
isAstral Char
c = Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
> Char
'\xFFFF'

utf16Width :: Char -> Int
utf16Width :: Char -> Int
utf16Width Char
c = if Char -> Bool
isAstral Char
c then Int
2 else Int
1

-- | The length of a text in UTF-16 code units.
utf16Length :: T.Text -> Int
utf16Length :: Text -> Int
utf16Length = (Int -> Char -> Int) -> Int -> Text -> Int
forall a. (a -> Char -> a) -> a -> Text -> a
T.foldl' (\Int
n Char
c -> Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Char -> Int
utf16Width Char
c) Int
0

-- | Convert a 0-based code-point column on the given 0-based line to UTF-16
-- units. A column beyond the end of the line keeps its distance past the end
-- (the diagnostics code uses column 99 to mean "to the end of the line").
colToUtf16 :: AstralLines -> Int -> Int -> Int
colToUtf16 :: AstralLines -> Int -> Int -> Int
colToUtf16 (AstralLines IntMap Text
ls) Int
line Int
col =
  case Int -> IntMap Text -> Maybe Text
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
line IntMap Text
ls of
    Maybe Text
Nothing   -> Int
col
    Just Text
text -> Text -> Int
utf16Length (Int -> Text -> Text
T.take Int
col Text
text) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
col Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
T.length Text
text)

-- | Convert a 0-based UTF-16 column on the given 0-based line to code
-- points. A column inside a surrogate pair maps to the start of its
-- character.
colFromUtf16 :: AstralLines -> Int -> Int -> Int
colFromUtf16 :: AstralLines -> Int -> Int -> Int
colFromUtf16 (AstralLines IntMap Text
ls) Int
line Int
col =
  case Int -> IntMap Text -> Maybe Text
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
line IntMap Text
ls of
    Maybe Text
Nothing   -> Int
col
    Just Text
text -> Int -> Int -> [Char] -> Int
go Int
0 Int
0 (Text -> [Char]
T.unpack Text
text)
      where
        go :: Int -> Int -> [Char] -> Int
go Int
cp Int
units (Char
c : [Char]
cs)
          | Int
units Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
col                = Int
cp
          | Int
units Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Char -> Int
utf16Width Char
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
col  = Int
cp
          | Bool
otherwise                   = Int -> Int -> [Char] -> Int
go (Int
cp Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
units Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Char -> Int
utf16Width Char
c) [Char]
cs
        go Int
cp Int
units []                  = Int
cp Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
col Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
units)

-- | Convert an incoming LSP position (UTF-16) to code points.
positionFromUtf16 :: AstralLines -> Position -> Position
positionFromUtf16 :: AstralLines -> Position -> Position
positionFromUtf16 AstralLines
als (Position UInt
l UInt
c) =
  UInt -> UInt -> Position
Position UInt
l (Int -> UInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral (AstralLines -> Int -> Int -> Int
colFromUtf16 AstralLines
als (UInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral UInt
l) (UInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral UInt
c)))

-- | Convert an outgoing range (code points) to UTF-16.
rangeToUtf16 :: AstralLines -> Range -> Range
rangeToUtf16 :: AstralLines -> Range -> Range
rangeToUtf16 AstralLines
als (Range Position
s Position
e) = Position -> Position -> Range
Range (Position -> Position
posToUtf16 Position
s) (Position -> Position
posToUtf16 Position
e)
  where
    posToUtf16 :: Position -> Position
posToUtf16 (Position UInt
l UInt
c) =
      UInt -> UInt -> Position
Position UInt
l (Int -> UInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral (AstralLines -> Int -> Int -> Int
colToUtf16 AstralLines
als (UInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral UInt
l) (UInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral UInt
c)))

-- | Convert semantic tokens (code points) to UTF-16. The length is converted
-- through the token's end column, so a token that itself contains astral
-- characters (e.g. the @𝕀@ keyword) gets its UTF-16 width.
tokensToUtf16 :: AstralLines -> [SemanticTokenAbsolute] -> [SemanticTokenAbsolute]
tokensToUtf16 :: AstralLines -> [SemanticTokenAbsolute] -> [SemanticTokenAbsolute]
tokensToUtf16 als :: AstralLines
als@(AstralLines IntMap Text
ls) [SemanticTokenAbsolute]
tokens
  | IntMap Text -> Bool
forall a. IntMap a -> Bool
IntMap.null IntMap Text
ls = [SemanticTokenAbsolute]
tokens
  | Bool
otherwise      = (SemanticTokenAbsolute -> SemanticTokenAbsolute)
-> [SemanticTokenAbsolute] -> [SemanticTokenAbsolute]
forall a b. (a -> b) -> [a] -> [b]
map SemanticTokenAbsolute -> SemanticTokenAbsolute
adjust [SemanticTokenAbsolute]
tokens
  where
    adjust :: SemanticTokenAbsolute -> SemanticTokenAbsolute
adjust SemanticTokenAbsolute
token = SemanticTokenAbsolute
token
      { _startChar = fromIntegral start'
      , _length    = fromIntegral (end' - start')
      }
      where
        line :: Int
line   = UInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SemanticTokenAbsolute -> UInt
_line SemanticTokenAbsolute
token)
        start :: Int
start  = UInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SemanticTokenAbsolute -> UInt
_startChar SemanticTokenAbsolute
token)
        start' :: Int
start' = AstralLines -> Int -> Int -> Int
colToUtf16 AstralLines
als Int
line Int
start
        end' :: Int
end'   = AstralLines -> Int -> Int -> Int
colToUtf16 AstralLines
als Int
line (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ UInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SemanticTokenAbsolute -> UInt
_length SemanticTokenAbsolute
token))