Fix bugs in lexeme tracking

8 files changed, 74 insertions, 44 deletions

diff --git a/src/dex.hs b/src/dex.hs
index 5232ec9c..2874649e 100644
--- a/src/dex.hs
+++ b/src/dex.hs
@@ -116,7 +116,7 @@ printFinal fmt prog = case fmt of
   TextDoc    -> return ()
   JSONDoc    -> return ()
 #ifdef DEX_LIVE
-  HTMLDoc    -> putStr $ progHtml prog
+  HTMLDoc    -> undefined -- putStr $ progHtml prog
 #endif
 
 readSourceBlock :: (MonadIO (m n), EnvReader m) => String -> m n SourceBlock
diff --git a/src/lib/ConcreteSyntax.hs b/src/lib/ConcreteSyntax.hs
index 2aa2c0f6..5762083e 100644
--- a/src/lib/ConcreteSyntax.hs
+++ b/src/lib/ConcreteSyntax.hs
@@ -197,12 +197,12 @@ topLevelCommand =
       importModule
   <|> declareForeign
   <|> declareCustomLinearization
-  <|> (Misc . QueryEnv <$> envQuery)
+  -- <|> (Misc . QueryEnv <$> envQuery)
   <|> explicitCommand
   <?> "top-level command"
 
-envQuery :: Parser EnvQuery
-envQuery = error "not implemented"
+_envQuery :: Parser EnvQuery
+_envQuery = error "not implemented"
 -- string ":debug" >> sc >> (
 --       (DumpSubst        <$  (string "env"   >> sc))
 --   <|> (InternalNameInfo <$> (string "iname" >> sc >> rawName))
diff --git a/src/lib/IncState.hs b/src/lib/IncState.hs
index 19d0a088..3c8f90d7 100644
--- a/src/lib/IncState.hs
+++ b/src/lib/IncState.hs
@@ -8,7 +8,7 @@
 
 module IncState (
   IncState (..), MapEltUpdate (..), MapUpdate (..),
-  Overwrite (..), TailUpdate (..)) where
+  Overwrite (..), TailUpdate (..), mapUpdateMapWithKey) where
 
 import qualified Data.Map.Strict as M
 import GHC.Generics
@@ -29,6 +29,9 @@ data MapEltUpdate v =
 data MapUpdate k v = MapUpdate { mapUpdates :: M.Map k (MapEltUpdate v) }
      deriving (Functor, Show, Generic)
 
+mapUpdateMapWithKey :: MapUpdate k v -> (k -> v -> v') -> MapUpdate k v'
+mapUpdateMapWithKey (MapUpdate m) f = MapUpdate $ M.mapWithKey (\k v -> fmap (f k) v) m
+
 instance Ord k => Monoid (MapUpdate k v) where
   mempty = MapUpdate mempty
 
diff --git a/src/lib/Lexing.hs b/src/lib/Lexing.hs
index d40c4c8a..18c85f55 100644
--- a/src/lib/Lexing.hs
+++ b/src/lib/Lexing.hs
@@ -235,12 +235,21 @@ space = gets canBreak >>= \case
   True  -> space1
   False -> void $ takeWhile1P (Just "white space") (`elem` (" \t" :: String))
 
+setCanBreakLocally :: Bool -> Parser a -> Parser a
+setCanBreakLocally brLocal p = do
+  brPrev <- gets canBreak
+  modify \ctx -> ctx {canBreak = brLocal}
+  ans <- p
+  modify \ctx -> ctx {canBreak = brPrev}
+  return ans
+{-# INLINE setCanBreakLocally #-}
+
 mayBreak :: Parser a -> Parser a
-mayBreak p = pLocal (\ctx -> ctx { canBreak = True }) p
+mayBreak p = setCanBreakLocally True p
 {-# INLINE mayBreak #-}
 
 mayNotBreak :: Parser a -> Parser a
-mayNotBreak p = pLocal (\ctx -> ctx { canBreak = False }) p
+mayNotBreak p = setCanBreakLocally False p
 {-# INLINE mayNotBreak #-}
 
 precededByWhitespace :: Parser Bool
@@ -294,14 +303,16 @@ withIndent p = do
   nextLine
   indent <- T.length <$> takeWhileP (Just "space") (==' ')
   when (indent <= 0) empty
-  pLocal (\ctx -> ctx { curIndent = curIndent ctx + indent }) $ mayNotBreak p
+  locallyExtendCurIndent indent $ mayNotBreak p
 {-# INLINE withIndent #-}
 
-pLocal :: (ParseCtx -> ParseCtx) -> Parser a -> Parser a
-pLocal f p = do
-  s <- get
-  put (f s) >> p <* put s
-{-# INLINE pLocal #-}
+locallyExtendCurIndent :: Int -> Parser a -> Parser a
+locallyExtendCurIndent n p = do
+  indentPrev <- gets curIndent
+  modify \ctx -> ctx { curIndent = indentPrev + n }
+  ans <- p
+  modify \ctx -> ctx { curIndent = indentPrev }
+  return ans
 
 eol :: Parser ()
 eol = void MC.eol
diff --git a/src/lib/Live/Eval.hs b/src/lib/Live/Eval.hs
index a8c449a8..f5248302 100644
--- a/src/lib/Live/Eval.hs
+++ b/src/lib/Live/Eval.hs
@@ -8,7 +8,7 @@
 {-# OPTIONS_GHC -Wno-orphans #-}
 
 module Live.Eval (
-  watchAndEvalFile, ResultsServer, ResultsUpdate, subscribeIO, dagAsUpdate) where
+  watchAndEvalFile, ResultsServer, ResultsUpdate, subscribeIO, dagAsUpdate, addSourceBlockIds) where
 
 import Control.Concurrent
 import Control.Monad
@@ -33,13 +33,16 @@ import MonadUtil
 
 -- `watchAndEvalFile` returns the channel by which a client may
 -- subscribe by sending a write-only view of its input channel.
-watchAndEvalFile :: FilePath -> EvalConfig -> TopStateEx
-                 -> IO (Evaluator SourceBlock Result)
+watchAndEvalFile :: FilePath -> EvalConfig -> TopStateEx -> IO ResultsServer
 watchAndEvalFile fname opts env = do
   watcher <- launchFileWatcher fname
   parser <- launchCellParser watcher \source -> uModuleSourceBlocks $ parseUModule Main source
   launchDagEvaluator parser env (evalSourceBlockIO opts)
 
+addSourceBlockIds :: NodeListUpdate (NodeState SourceBlock o) -> NodeListUpdate (NodeState SourceBlockWithId o)
+addSourceBlockIds (NodeListUpdate listUpdate mapUpdate) = NodeListUpdate listUpdate mapUpdate'
+   where mapUpdate' = mapUpdateMapWithKey mapUpdate \k (NodeState i o) -> NodeState (SourceBlockWithId k i) o
+
 type ResultsServer = Evaluator        SourceBlock Result
 type ResultsUpdate = EvalStatusUpdate SourceBlock Result
 
@@ -297,15 +300,15 @@ processDagUpdate dagUpdate = do
 
 -- === instances ===
 
-instance ToJSON a => ToJSON (NodeListUpdate a)
+instance (ToJSON i, ToJSON o) => ToJSON (NodeListUpdate (NodeState i o)) where
 instance (ToJSON a, ToJSONKey k) => ToJSON (MapUpdate k a)
 instance ToJSON a => ToJSON (TailUpdate a)
 instance ToJSON a => ToJSON (MapEltUpdate a)
 instance ToJSON o => ToJSON (NodeEvalStatus o)
 instance (ToJSON i, ToJSON o) => ToJSON (NodeState i o)
 
-instance ToJSON SourceBlock where
-  toJSON b = toJSON (sbLine b, pprintHtml b)
+instance ToJSON SourceBlockWithId where
+  toJSON b@(SourceBlockWithId _ b') = toJSON (sbLine b', pprintHtml b)
 instance ToJSON Result      where toJSON = toJSONViaHtml
 
 toJSONViaHtml :: ToMarkup a => a -> Value
diff --git a/src/lib/Live/Web.hs b/src/lib/Live/Web.hs
index ae727302..1ab7394c 100644
--- a/src/lib/Live/Web.hs
+++ b/src/lib/Live/Web.hs
@@ -57,7 +57,7 @@ resultStream resultsServer write flush = do
   where
     sendUpdate :: ResultsUpdate -> IO ()
     sendUpdate update = do
-      let s = encodeResults update
+      let s = encodeResults $ addSourceBlockIds update
       write (fromByteString s) >> flush
 
     encodeResults :: ToJSON a => a -> BS.ByteString
diff --git a/src/lib/RenderHtml.hs b/src/lib/RenderHtml.hs
index fb9e7fad..7e4aac45 100644
--- a/src/lib/RenderHtml.hs
+++ b/src/lib/RenderHtml.hs
@@ -16,6 +16,7 @@ import Text.Blaze.Html.Renderer.String
 import qualified Data.Map.Strict as M
 import Control.Monad.State.Strict
 import Data.Maybe (fromJust)
+import Data.String (fromString)
 import Data.Text    qualified as T
 import Data.Text.IO qualified as T
 import CMark (commonmarkToHtml)
@@ -72,10 +73,10 @@ instance ToMarkup Output where
     HtmlOut s -> preEscapedString s
     _ -> cdiv "result-block" $ toHtml $ pprint out
 
-instance ToMarkup SourceBlock where
-  toMarkup block = case sbContents block of
-    (Misc (ProseBlock s)) -> cdiv "prose-block" $ mdToHtml s
-    _ -> renderSpans (sbLexemeInfo block) (sbASTInfo block) (sbText block)
+instance ToMarkup SourceBlockWithId where
+  toMarkup (SourceBlockWithId blockId block) = case sbContents block of
+    Misc (ProseBlock s) -> cdiv "prose-block" $ mdToHtml s
+    _ -> renderSpans blockId (sbLexemeInfo block) (sbASTInfo block) (sbText block)
 
 mdToHtml :: T.Text -> Html
 mdToHtml s = preEscapedText $ commonmarkToHtml [] s
@@ -83,28 +84,37 @@ mdToHtml s = preEscapedText $ commonmarkToHtml [] s
 cdiv :: String -> Html -> Html
 cdiv c inner = H.div inner ! class_ (stringValue c)
 
-renderSpans :: LexemeInfo -> ASTInfo -> T.Text -> Markup
-renderSpans lexInfo astInfo sourceText = cdiv "code-block" do
+type BlockId = Int
+
+renderSpans :: BlockId -> LexemeInfo -> ASTInfo -> T.Text -> Markup
+renderSpans blockId lexInfo astInfo sourceText = cdiv "code-block" do
   runTextWalkerT sourceText do
     forM_ (lexemeList lexInfo) \sourceId -> do
       let (lexemeTy, (l, r)) = fromJust $ M.lookup sourceId (lexemeInfo lexInfo)
-      takeTo l >>= emitSpan ""
-      takeTo r >>= emitSpan (lexemeClass lexemeTy)
-    takeRest >>= emitSpan ""
-
-emitSpan :: String -> T.Text -> TextWalker ()
-emitSpan className t = lift $ H.span (toHtml t) ! class_ (stringValue className)
-
-lexemeClass :: LexemeType -> String
+      takeTo l >>= emitSpan Nothing (Just "comment")
+      takeTo r >>= emitSpan (Just (blockId, sourceId)) (lexemeClass lexemeTy)
+    takeRest >>= emitSpan Nothing (Just "comment")
+
+emitSpan :: Maybe (BlockId, SrcId) -> Maybe String -> T.Text -> TextWalker ()
+emitSpan maybeSrcId className t = lift do
+  let classAttr = case className of
+        Nothing -> mempty
+        Just c -> class_ (stringValue c)
+  let idAttr = case maybeSrcId of
+        Nothing -> mempty
+        Just (bid, SrcId sid) -> At.id (fromString $ "span_" ++ show bid ++ "_"++ show sid)
+  H.span (toHtml t) ! classAttr ! idAttr
+
+lexemeClass :: LexemeType -> Maybe String
 lexemeClass = \case
-   Keyword             -> "keyword"
-   Symbol              -> "symbol"
-   TypeName            -> "type-name"
-   LowerName           -> ""
-   UpperName           -> ""
-   LiteralLexeme       -> "literal"
-   StringLiteralLexeme -> ""
-   MiscLexeme          -> ""
+   Keyword             -> Just "keyword"
+   Symbol              -> Just "symbol"
+   TypeName            -> Just "type-name"
+   LowerName           -> Nothing
+   UpperName           -> Nothing
+   LiteralLexeme       -> Just "literal"
+   StringLiteralLexeme -> Nothing
+   MiscLexeme          -> Nothing
 
 type TextWalker a = StateT (Int, T.Text) MarkupM a
 
@@ -121,5 +131,6 @@ takeTo startPos = do
 
 takeRest :: TextWalker T.Text
 takeRest = do
-  endPos <- gets $ T.length . snd
-  takeTo endPos
+  (curPos, curText) <- get
+  put (curPos + T.length curText, mempty)
+  return curText
diff --git a/src/lib/Types/Source.hs b/src/lib/Types/Source.hs
index 961445ea..81f1c338 100644
--- a/src/lib/Types/Source.hs
+++ b/src/lib/Types/Source.hs
@@ -509,6 +509,8 @@ data UModule = UModule
 
 -- === top-level blocks ===
 
+data SourceBlockWithId = SourceBlockWithId Int SourceBlock
+
 data SourceBlock = SourceBlock
   { sbLine       :: Int
   , sbOffset     :: Int