Reader.hs

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

module FNotation.Reader where

import Data.IORef
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Text (Text)
import Data.Vector qualified as V
import Diagnostician
import FNotation.Config
import FNotation.Kinds qualified as T
import FNotation.Names
import FNotation.Tokens qualified as T
import FNotation.Trees hiding (head)
import Prettyprinter
import Prelude hiding (head, lookup)

-- Reader diagnostics
--------------------------------------------------------------------------------

data ReaderCode
  = UnexpectedToken
  | DefaultedPrec
  | IncompatiblePrecedences
  | ModifierWithoutModifyee
  | StatementWithoutNewline
  deriving (Eq, Ord)

readerCodeTable :: Map ReaderCode CodeMeta
readerCodeTable =
  Map.fromList
    [ (UnexpectedToken, CodeMeta 0 SError Nothing)
    , (DefaultedPrec, CodeMeta 1 SWarning Nothing)
    , (IncompatiblePrecedences, CodeMeta 2 SError Nothing)
    , (ModifierWithoutModifyee, CodeMeta 3 SError Nothing)
    , (StatementWithoutNewline, CodeMeta 4 SError Nothing)
    ]

-- Reader monad
--------------------------------------------------------------------------------

data ReaderState = ReaderState
  { pos :: IORef Int
  , gas :: IORef Int
  , skipNewlines :: IORef Bool
  , tokens :: V.Vector T.Token
  , file :: File
  , reporter :: Reporter ReaderCode
  , config :: ConfTable Prec
  }

-- Parsing utilities
--------------------------------------------------------------------------------

report :: ReaderState -> Span -> ReaderCode -> DDoc -> IO ()
report st s c m = do
  let n = Note (Just (SourceLoc st.file s)) Nothing
  let d = Diagnostic c m [n]
  st.reporter.reportIO d

cur :: ReaderState -> IO T.Kind
cur st = do
  gas <- readIORef st.gas
  if gas <= 0
    then do
      pos <- readIORef st.pos
      let token = st.tokens V.! pos
      error $ "out of gas at token " ++ show (dpretty token)
    else writeIORef st.gas (gas - 1)
  pos <- readIORef st.pos
  pure (st.tokens V.! pos).kind

locally :: IORef a -> a -> IO b -> IO b
locally ref v action = do
  old <- readIORef ref
  writeIORef ref v
  res <- action
  writeIORef ref old
  pure res

ignoreNewlines :: ReaderState -> IO a -> IO a
ignoreNewlines st = locally st.skipNewlines True

withNewlines :: ReaderState -> IO a -> IO a
withNewlines st = locally st.skipNewlines False

curSpan :: ReaderState -> IO Span
curSpan st = do
  pos <- readIORef st.pos
  pure (V.unsafeIndex st.tokens pos).span

curValue :: ReaderState -> IO T.TokenValue
curValue st = do
  pos <- readIORef st.pos
  pure (V.unsafeIndex st.tokens pos).value

curName :: ReaderState -> IO Name
curName st =
  curValue st >>= \case
    T.VName x -> pure x
    _ -> error "expected token to be associated with a name"

curInt :: ReaderState -> IO Int
curInt st =
  curValue st >>= \case
    T.VInt x -> pure x
    _ -> error "expected token to be associated with an int"

curString :: ReaderState -> IO Text
curString st =
  curValue st >>= \case
    T.VString x -> pure x
    _ -> error "expected token to be associated with a string"

at :: ReaderState -> T.Kind -> IO Bool
at st k = (k ==) <$> cur st

advance :: ReaderState -> IO ()
advance st = do
  pos <- readIORef st.pos
  let n = V.length st.tokens
  if pos < n - 1
    then do
      let next j
            | j < n = case (V.unsafeIndex st.tokens j).kind of
                T.Nl -> next (j + 1)
                _ -> j
            | otherwise = j
      readIORef st.skipNewlines >>= \case
        True -> writeIORef st.pos $ next (pos + 1)
        False -> writeIORef st.pos $ pos + 1
      writeIORef st.gas 256
    else pure ()

eat :: ReaderState -> T.Kind -> IO Bool
eat st k =
  at st k >>= \case
    True -> advance st >> pure True
    False -> pure False

reportUnexpected :: ReaderState -> T.Kind -> T.Class -> IO ()
reportUnexpected st k c = do
  s <- curSpan st
  report st s UnexpectedToken $
    "Unexpected token kind" <+> dpretty k <> ", expected" <+> dpretty c

expect :: ReaderState -> T.Kind -> IO ()
expect st k = do
  k' <- cur st
  if k == k'
    then advance st
    else
      reportUnexpected st k' (T.CSpecific k) >> pure ()

openingPos :: ReaderState -> IO Pos
openingPos st = (.start) <$> curSpan st

close :: ReaderState -> Pos -> (Span -> Ntn) -> IO Ntn
close st s f = do
  (Span _ e) <- curSpan st
  pure $ f (Span s e)

advanceClose :: ReaderState -> Pos -> (Span -> Ntn) -> IO Ntn
advanceClose st s f = do
  n <- close st s f
  advance st
  pure n

-- The fnotation grammar
--------------------------------------------------------------------------------

argStarts :: V.Vector T.Kind
argStarts =
  V.fromList
    [ T.LParen
    , T.LBrack
    , T.AIdent
    , T.AKeyword
    , T.Field
    , T.Tag
    , T.Int
    , T.Block
    ]

argStart :: T.Kind -> Bool
argStart k = V.elem k argStarts

tupleElems :: ReaderState -> IO [Ntn]
tupleElems st =
  cur st >>= \case
    T.RBrack -> pure []
    k | argStart k -> do
      n <- expr st
      cur st >>= \case
        T.RBrack -> pure [n]
        T.Comma -> do
          advance st
          ns <- tupleElems st
          pure $ n : ns
        k' -> do
          reportUnexpected st k' T.CTupleMark
          pure [n]
    k -> do
      reportUnexpected st k T.CExprStart
      pure []

argBase :: ReaderState -> IO Ntn
argBase st = do
  m <- openingPos st
  cur st >>= \case
    T.LParen -> do
      e <- ignoreNewlines st $ do
        advance st
        expr st
      expect st T.RParen
      pure e
    T.LBrack -> do
      ns <- ignoreNewlines st $ do
        advance st
        tupleElems st
      expect st T.RBrack
      close st m $ Tuple ns
    T.AIdent -> do
      x <- curName st
      advanceClose st m $ Ident x
    T.AKeyword -> do
      x <- curName st
      advanceClose st m $ Keyword x
    T.Field; T.FieldImmediate -> do
      x <- curName st
      advanceClose st m $ Field x
    T.Tag -> do
      x <- curName st
      advanceClose st m $ Tag x
    T.Int -> do
      i <- curInt st
      advanceClose st m $ Int i
    T.String -> do
      x <- curString st
      advanceClose st m $ String x
    T.Block -> block st
    k -> do
      reportUnexpected st k T.CExprStart
      advanceClose st m Error

-- `.x.y.z -> [x, y, z]`
argProjs :: ReaderState -> IO [Ntn]
argProjs st =
  cur st >>= \case
    T.FieldImmediate -> do
      field <- argBase st
      rest <- argProjs st
      return (field : rest)
    _ -> pure []

arg :: ReaderState -> IO Ntn
arg st = do
  head <- argBase st
  spine <- argProjs st
  pure $ foldr (flip Juxt) head (reverse spine)

args :: ReaderState -> IO [Ntn]
args st = do
  k <- cur st
  if argStart k
    then (:) <$> arg st <*> args st
    else pure []

expr :: ReaderState -> IO Ntn
expr st = arg st >>= go (Prec 0 AssocNon)
 where
  go p lhs = do
    cur st >>= \case
      k@(T.SIdent; T.SKeyword) -> do
        s <- curSpan st
        x <- curName st
        let n = case k of
              T.SIdent -> Ident x s
              T.SKeyword -> Keyword x s
        p' <- case confTableLookup st.config x.last of
          Just p' -> pure p'
          Nothing -> do
            report st s DefaultedPrec $
              "Defaulted precedence of" <+> dpretty x <+> "to the same as +"
            pure $ Prec 50 AssocL
        case precLe p p' of
          Nothing -> do
            report st s IncompatiblePrecedences "Incompatible precedences"
            pure lhs
          Just False -> pure lhs
          Just True -> do
            advance st
            rhs <- arg st >>= go p'
            go p (Infix lhs n rhs)
      k | argStart k -> do
        rhs <- arg st
        go p (Juxt lhs rhs)
      _ -> pure lhs

decl :: ReaderState -> IO Ntn
decl st = do
  p <- openingPos st
  go p []
 where
  go p mods =
    cur st >>= \case
      T.Modifier -> do
        m <- curName st
        advance st
        go p (m : mods)
      T.Decl -> do
        x <- curName st
        advance st
        n <- expr st
        pure $ MDecl (reverse mods) x n (Span p (endPos n))
      _ -> do
        s <- curSpan st
        report st s ModifierWithoutModifyee "expected a declaration after a declaration modifier"
        advanceClose st p Error

stmt :: ReaderState -> IO Ntn
stmt st =
  cur st >>= \case
    T.Modifier; T.Decl -> decl st
    _ -> expr st

stmts :: ReaderState -> IO [Ntn]
stmts st = go True []
 where
  -- following = we have seen at least one newline
  go following ns =
    cur st >>= \case
      T.Nl -> do
        advance st
        go True ns
      k | k == T.End || k == T.Eof -> pure $ reverse ns
      _ -> case following of
        False -> do
          s <- curSpan st
          report st s StatementWithoutNewline "expected a newline, end, or eof after a statement"
          pure $ reverse ns
        True -> do
          n <- stmt st
          go False $ n : ns

block :: ReaderState -> IO Ntn
block st =
  cur st >>= \case
    T.Block -> do
      m <- openingPos st
      x <- curName st
      advance st
      h <-
        cur st >>= \case
          k | argStart k -> Just <$> expr st
          _ -> pure Nothing
      ns <- stmts st
      advanceClose st m $ Block x h ns
    _ -> error "expected a block"

-- Toplevel parsing interface
--------------------------------------------------------------------------------

read :: ConfTable Prec -> Reporter ReaderCode -> File -> V.Vector T.Token -> IO [Ntn]
read config reporter file tokens = do
  pos <- newIORef 0
  gas <- newIORef 256
  skipNewlines <- newIORef False
  let st = ReaderState pos gas skipNewlines tokens file reporter config
  stmts st