diagnostician-0.2.0.0: src/Diagnostician.hs

-- SPDX-FileCopyrightText: 2026 Coln contributors
--
-- SPDX-License-Identifier: Apache-2.0 OR MIT

module Diagnostician where

import Data.Functor.Contravariant
import Data.IORef (IORef, atomicModifyIORef)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (maybeToList)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Unsafe qualified as TU
import Data.Vector.Unboxed qualified as UV
import Prettyprinter
import Prettyprinter.Render.Text
import System.IO (Handle)

-- Pretty printer annotations
--------------------------------------------------------------------------------

data DiagnosticAnn
  = DSeverity
  | DCode
  | DBar
  | DSpan

-- Diagnostician doc
type DDoc = Doc DiagnosticAnn

-- | Pretty for diagnostics
class DPretty a where
  dpretty :: a -> DDoc

-- Source locations
--------------------------------------------------------------------------------

type Pos = Int

data Span = Span {start :: Int, end :: Int}
  deriving (Eq)

instance DPretty Span where
  dpretty (Span s e) = pretty s <> ":" <> pretty e

-- Util
--------------------------------------------------------------------------------

sliceWord8 :: Pos -> Pos -> Text -> Text
sliceWord8 s e t = TU.dropWord8 s $ TU.takeWord8 e t

-- Files
--------------------------------------------------------------------------------

{- | A @File@ is used to display diagnostic messages.
Specifically, a @File@ is used to convert the @Span@ in a diagnostic message
into snippet of source code with the span underlined.

In order to do this, we need to convert the byte positions in the @Span@ into
line/column positions. This can be done fairly efficiently by binary search
through the vector of newline positions in the file, so we create this vector
whenever we open a file and store it in the @File@ record.
-}
data File = File
  { name :: FilePath
  , contents :: T.Text
  , lineBreaks :: UV.Vector Pos
  }

newFile :: FilePath -> T.Text -> File
newFile x t = File x t (findLineBreaks t)

findLineBreaks :: T.Text -> UV.Vector Pos
findLineBreaks t = UV.unfoldr nextNewline (-1)
 where
  l = TU.lengthWord8 t
  nextNewline i
    | i == -1 = Just (-1, 0)
    | i < l = case TU.iter t i of
        TU.Iter '\n' j -> Just (i, i + j)
        TU.Iter _ j -> nextNewline (i + j)
    | i == l = Just (i, i + 1)
    | otherwise = Nothing

type LineNum = Int

type ColNum = Int

lineStart :: File -> LineNum -> Pos
lineStart f l = (f.lineBreaks UV.! l) + 1

lineEnd :: File -> LineNum -> Pos
lineEnd f l = f.lineBreaks UV.! (l + 1)

lineSpan :: File -> LineNum -> Span
lineSpan f l = Span (lineStart f l) (lineEnd f l)

lineContents :: File -> LineNum -> T.Text
lineContents f l = sliceWord8 (lineStart f l) (lineEnd f l) f.contents

srcOf :: File -> Pos -> (LineNum, ColNum)
srcOf f i =
  seq
    (0 <= i && i <= TU.lengthWord8 f.contents || error "position out of bounds")
    (go 0 (UV.length f.lineBreaks - 1))
 where
  go l r
    | l == r = (l, c l)
    | i < lineStart f m = go l m
    | i > lineEnd f m = go m r
    | otherwise = (m, c m)
   where
    m = (l + r) `div` 2
    c lineNo = i - lineStart f lineNo

repeated :: Int -> Char -> Doc ann
repeated n c
  | n == 0 = mempty
  | n == 1 = pretty c
  | otherwise = pretty (T.replicate n (T.singleton c))

linePretty :: Int -> LineNum -> Span -> T.Text -> Span -> DDoc
linePretty numWidth l (Span ls le) t (Span s e) =
  vsep
    [ gutter <+> pretty t
    , gutter <+> repeated ns ' ' <> annotate DSpan (repeated nc '^')
    ]
 where
  s' = max ls s
  e' = min le e
  ns = T.length $ sliceWord8 0 (s' - ls) t
  nc = max 1 $ T.length $ sliceWord8 (s' - ls) (e' - ls) t
  ln = fill numWidth $ pretty $ l + 1
  gutter = ln <+> annotate DBar "|"

numDigits :: Int -> Int
numDigits n = go (abs n)
 where
  go x
    | x < 10 = 1
    | otherwise = 1 + go (x `div` 10)

-- | This is the function used to display the source code for a @Span@.
linesPretty :: File -> Span -> DDoc
linesPretty f sp@(Span s e) =
  vsep
    [ linePretty numWidth l (lineSpan f l) (lineContents f l) sp
    | l <- [ls .. le]
    ]
 where
  ls = fst $ srcOf f s
  le = fst $ srcOf f e
  numWidth = max (numDigits ls) (numDigits le)

-- Reporter
--------------------------------------------------------------------------------

{- | A @Reporter@ is a destination for diagnostic messages. It corresponds
to an IO action `reportIO` that logs a diagnostic
-}
newtype Reporter a = Reporter {reportIO :: Diagnostic a -> IO ()}

reportTo :: Reporter a -> Diagnostic a -> IO ()
reportTo r d = r.reportIO d

instance Contravariant Reporter where
  contramap f (Reporter r) = Reporter $ r . fmap f

-- | Create a `Reporter` that writes the prettyprinted diagnostic to a file
fileReporter :: (Code a) => Handle -> Reporter a
fileReporter handle =
  Reporter
    { reportIO = \d -> hPutDoc handle (hardline <> dpretty d <> hardline)
    }

-- | Create a `Reporter` that appends a diagnostic to a list
pureReporter :: (Code a) => IORef [Diagnostic a] -> Reporter a
pureReporter ref =
  Reporter
    { reportIO = \d -> atomicModifyIORef ref (\ds -> (d : ds, ()))
    }

-- Diagnostics
--------------------------------------------------------------------------------

data Severity = SDebug | SInfo | SWarning | SError

data CodeMeta = CodeMeta
  { number :: Int
  , severity :: Severity
  , about :: Maybe Text
  }

class Code a where
  codeMeta :: a -> CodeMeta

promoteCodeTable :: (Ord b) => Map a CodeMeta -> (a -> b) -> Int -> Map b CodeMeta
promoteCodeTable t f offset =
  Map.fromList
    [(f c, m{number = m.number + offset}) | (c, m) <- Map.toList t]

padWithZerosTo :: Int -> Int -> Doc ann
padWithZerosTo w i = repeated (w - numDigits i) '0' <> pretty i

prtCode :: (Code a) => a -> DDoc
prtCode c = annotate DSeverity s <> "[" <> annotate DCode (sl <> padWithZerosTo 4 m.number) <> "]"
 where
  m = codeMeta c
  (s, sl) = case m.severity of
    SDebug -> ("debug", "D")
    SInfo -> ("info", "I")
    SWarning -> ("warning", "W")
    SError -> ("error", "E")

data SourceLoc = SourceLoc
  { file :: File
  , span :: Span
  }

instance DPretty SourceLoc where
  dpretty (SourceLoc f s) = linesPretty f s

data Note = Note
  { noteSourceLoc :: Maybe SourceLoc
  , noteMessage :: Maybe DDoc
  }

instance DPretty Note where
  dpretty (Note loc message) =
    vsep $
      (dpretty <$> maybeToList loc) ++ (unAnnotate <$> maybeToList message)

data Diagnostic a = Diagnostic
  { code :: a
  , summary :: DDoc
  , notes :: [Note]
  }
  deriving (Functor)

instance (Code a) => DPretty (Diagnostic a) where
  dpretty d =
    vsep $
      (prtCode d.code <> ": " <> unAnnotate d.summary) : (map dpretty d.notes)