Organize source information by pass and do more of the work on the Haskell side.

This way we keep the serialized data type more stable since that's the only part
of the system that isn't type checked.

12 files changed, 276 insertions, 127 deletions

diff --git a/src/dex.hs b/src/dex.hs
index 8c130d85..f16aa12a 100644
--- a/src/dex.hs
+++ b/src/dex.hs
@@ -192,7 +192,7 @@ parseEvalOpts = EvalConfig
                 , ("debug" , DebugLogLevel ) ]
 
 stdOutLogger :: Outputs -> IO ()
-stdOutLogger outs = do
+stdOutLogger (Outputs outs) = do
   isatty <- queryTerminal stdOutput
   forM_ outs \out -> putStr $ printOutput isatty out
 
diff --git a/src/lib/IncState.hs b/src/lib/IncState.hs
index 170a5b0d..a9a2bcb4 100644
--- a/src/lib/IncState.hs
+++ b/src/lib/IncState.hs
@@ -27,7 +27,7 @@ data MapEltUpdate s d =
  | Replace s  -- TODO: should we merge Create/Replace?
  | Update d
  | Delete
- deriving (Functor, Show, Generic)
+ deriving (Eq, Functor, Show, Generic)
 
 newtype MapUpdate k s d = MapUpdate { mapUpdates :: M.Map k (MapEltUpdate s d) }
      deriving (Functor, Show, Generic)
diff --git a/src/lib/Inference.hs b/src/lib/Inference.hs
index e3687c34..fec4a2fb 100644
--- a/src/lib/Inference.hs
+++ b/src/lib/Inference.hs
@@ -49,11 +49,11 @@ import Util hiding (group)
 
 -- === Top-level interface ===
 
-checkTopUType :: (Fallible1 m, TopLogger m, EnvReader m) => UType n -> m n (CType n)
+checkTopUType :: (Fallible1 m, TopLogger1 m, EnvReader m) => UType n -> m n (CType n)
 checkTopUType ty = liftInfererM $ checkUType ty
 {-# SCC checkTopUType #-}
 
-inferTopUExpr :: (Fallible1 m, TopLogger m, EnvReader m) => UExpr n -> m n (TopBlock CoreIR n)
+inferTopUExpr :: (Fallible1 m, TopLogger1 m, EnvReader m) => UExpr n -> m n (TopBlock CoreIR n)
 inferTopUExpr e = fst <$> (asTopBlock =<< liftInfererM (buildBlock $ bottomUp e))
 {-# SCC inferTopUExpr #-}
 
@@ -62,9 +62,7 @@ data UDeclInferenceResult e n =
  | UDeclResultBindName LetAnn (TopBlock CoreIR n) (Abs (UBinder (AtomNameC CoreIR)) e n)
  | UDeclResultBindPattern NameHint (TopBlock CoreIR n) (ReconAbs CoreIR e n)
 
-type TopLogger (m::MonadKind1) = forall n. Logger Outputs (m n)
-
-inferTopUDecl :: (Mut n, Fallible1 m, TopBuilder m, HasNamesE e, TopLogger m)
+inferTopUDecl :: (Mut n, Fallible1 m, TopBuilder m, HasNamesE e, TopLogger1 m)
               => UTopDecl n l -> e l -> m n (UDeclInferenceResult e n)
 inferTopUDecl (UStructDecl tc def) result = do
   tc' <- emitBinding (getNameHint tc) $ TyConBinding Nothing (DotMethods mempty)
@@ -165,10 +163,10 @@ newtype InfererM (i::S) (o::S) (a:: *) = InfererM
 type InfererCPSB  b i    o a = (forall o'. DExt o o' => b o o' -> InfererM i  o' a) -> InfererM i o a
 type InfererCPSB2 b i i' o a = (forall o'. DExt o o' => b o o' -> InfererM i' o' a) -> InfererM i o a
 
-liftInfererM :: (EnvReader m, TopLogger m, Fallible1 m) => InfererM n n a -> m n a
+liftInfererM :: (EnvReader m, TopLogger1 m, Fallible1 m) => InfererM n n a -> m n a
 liftInfererM cont = do
   (ansExcept, typeInfo) <- liftInfererMPure cont
-  emitLog [SourceInfo $ SITypeInfo typeInfo]
+  emitLog $ Outputs [SourceInfo $ SITypeInfo typeInfo]
   liftExcept ansExcept
 {-# INLINE liftInfererM #-}
 
@@ -350,9 +348,7 @@ withAllowedEffects effs cont = withInfState (\(InfState g _) -> InfState g effs)
 {-# INLINE withAllowedEffects #-}
 
 emitTypeInfo :: SrcId -> String -> InfererM i o ()
-emitTypeInfo sid ty = do
-  InfererM $ liftSubstReaderT $ lift11 $ lift1 $ lift do
-    modify \(TypeInfo m) -> TypeInfo $ M.insert sid ty m
+emitTypeInfo _ _ = return ()
 
 withReducibleEmissions
   :: (HasNamesE e, SubstE AtomSubstVal e, ToErr err)
diff --git a/src/lib/LLVM/Compile.hs b/src/lib/LLVM/Compile.hs
index 6a306fa3..436495ff 100644
--- a/src/lib/LLVM/Compile.hs
+++ b/src/lib/LLVM/Compile.hs
@@ -109,7 +109,7 @@ standardCompilationPipeline opt logger exports tm m = do
       s <- case ioLogLevel logger of
         DebugLogLevel -> Just <$> showIt
         NormalLogLevel -> return Nothing
-      ioLogAction logger [PassResult passName s]
+      ioLogAction logger $ Outputs [PassResult passName s]
 {-# SCC standardCompilationPipeline #-}
 
 internalize :: [String] -> Mod.Module -> IO ()
diff --git a/src/lib/RenderHtml.hs b/src/lib/RenderHtml.hs
index 4b87f96e..4015cfb7 100644
--- a/src/lib/RenderHtml.hs
+++ b/src/lib/RenderHtml.hs
@@ -8,7 +8,7 @@
 {-# OPTIONS_GHC -Wno-incomplete-patterns #-}
 
 module RenderHtml (
-  progHtml, pprintHtml, ToMarkup, renderSourceBlock, renderOutputs,
+  progHtml, pprintHtml, ToMarkup, renderSourceBlock,
   RenderedSourceBlock, RenderedOutputs, SourceBlockWithId (..)) where
 
 import Text.Blaze.Internal (MarkupM)
@@ -18,7 +18,9 @@ import Text.Blaze.Html.Renderer.String
 import Data.Aeson (ToJSON (..))
 import qualified Data.Map.Strict as M
 import Control.Monad.State.Strict
+import Data.Foldable (fold)
 import Data.Maybe (fromJust)
+import Data.Functor ((<&>))
 import Data.String (fromString)
 import Data.Text    qualified as T
 import Data.Text.IO qualified as T
@@ -31,6 +33,7 @@ import Paths_dex (getDataFileName)
 import PPrint
 import Types.Source
 import Util (unsnoc)
+import IncState
 
 -- === rendering source blocks and results ===
 
@@ -43,8 +46,8 @@ data SourceBlockWithId = SourceBlockWithId
 instance ToJSON SourceBlockWithId where
   toJSON (SourceBlockWithId n b) = toJSON $ renderSourceBlock n b
 
-instance ToJSON Output where
-  toJSON x = toJSON $ renderOutput x
+instance ToJSON Outputs where
+  toJSON x = toJSON $ renderOutputs x
 
 -- === rendering results ===
 
@@ -62,30 +65,51 @@ data RenderedSourceBlock = RenderedSourceBlock
   deriving (Generic)
 
 type RenderedOutputs = [RenderedOutput]
--- This is extremely close to `Output` and we could just use that directly. The
--- reason to keep it separate is that the haskell-javascript boundary is a bit
--- delicate and this provides some robustness against future changes to
--- `Output`.
 data RenderedOutput =
     RenderedTextOut String
   | RenderedHtmlOut String
-  | RenderedSourceInfo SourceInfo  -- for hovertips etc
   | RenderedPassResult PassName (Maybe String)
   | RenderedMiscLog String
   | RenderedError (Maybe SrcId) String
+  | RenderedTreeNodeUpdate [(SrcId, MapEltUpdate TreeNodeState TreeNodeUpdate)]
+  | RenderedFocusUpdate [(LexemeId, SrcId)]
+    deriving (Show, Eq, Generic)
+
+data HighlightType =
+   HighlightGroup
+ | HighlightLeaf
+ | HighlightError
+ | HighlightBinder
+ | HighlightOcc
+ | HighlightScope
+   deriving (Show, Eq, Generic)
+
+type RenderedHighlight = (SrcId, HighlightType)
+data TreeNodeState  = TreeNodeState
+  { tnSpan       :: (LexemeId, LexemeId)
+  , tnHighlights :: [RenderedHighlight]
+  , tnText       :: String }
+  deriving (Show, Eq, Generic)
+
+data TreeNodeUpdate = TreeNodeUpdate
+  { tnuHighlights :: Overwrite [RenderedHighlight]
+  , tnuText       :: Overwrite String }
     deriving (Show, Eq, Generic)
 
 renderOutputs :: Outputs -> RenderedOutputs
-renderOutputs outs = map renderOutput outs
+renderOutputs (Outputs outs) = foldMap renderOutput outs
 
-renderOutput :: Output -> RenderedOutput
+renderOutput :: Output -> [RenderedOutput]
 renderOutput = \case
-  TextOut s      -> RenderedTextOut s
-  HtmlOut s      -> RenderedHtmlOut s
-  SourceInfo si  -> RenderedSourceInfo si
-  PassResult n s -> RenderedPassResult n s
-  MiscLog s      -> RenderedMiscLog s
-  Error e        -> RenderedError (getErrSrcId e) (pprint e)
+  TextOut s      -> pure $ RenderedTextOut s
+  HtmlOut s      -> pure $ RenderedHtmlOut s
+  SourceInfo s   -> case s of
+    SIGroupingInfo  info -> renderGroupingInfo  info
+    SINamingInfo    info -> renderNamingInfo    info
+    SITypeInfo      info -> renderTypeInfo      info
+  PassResult n s -> pure $ RenderedPassResult n s
+  MiscLog s      -> pure $ RenderedMiscLog s
+  Error e        -> pure $ RenderedError (getErrSrcId e) (pprint e)
 
 renderSourceBlock :: BlockId -> SourceBlock -> RenderedSourceBlock
 renderSourceBlock n b = RenderedSourceBlock
@@ -99,8 +123,48 @@ renderSourceBlock n b = RenderedSourceBlock
       _ -> renderSpans n (sbLexemeInfo b) (sbText b)
   }
 
+renderGroupingInfo :: GroupingInfo -> RenderedOutputs
+renderGroupingInfo (GroupingInfo m) =
+  [ RenderedFocusUpdate    focus
+  , RenderedTreeNodeUpdate treeNodeUpdate]
+  where
+    treeNodeUpdate = M.toList m <&> \(sid, node) -> (sid, Create $ renderTreeNode m sid node)
+    focus = fold $ uncurry renderFocus <$> M.toList m
+
+renderTreeNode :: M.Map SrcId GroupTreeNode -> SrcId -> GroupTreeNode -> TreeNodeState
+renderTreeNode m sid node = TreeNodeState (gtnSpan node) (getHighlights m sid node) ""
+
+getHighlights :: M.Map SrcId GroupTreeNode -> SrcId -> GroupTreeNode -> [RenderedHighlight]
+getHighlights m sid node = case gtnChildren node of
+  [] -> [(sid, HighlightGroup)]
+  children -> children <&> \childSrcId -> do
+    let child = fromJust $ M.lookup childSrcId m
+    let highlight = case gtnIsAtomicLexeme child of
+                      True  -> HighlightLeaf
+                      False -> HighlightGroup
+    (childSrcId, highlight)
+
+renderFocus :: SrcId -> GroupTreeNode -> [(LexemeId, SrcId)]
+renderFocus srcId node = case gtnChildren node of
+  [] -> case gtnIsAtomicLexeme node of
+    False -> [(srcId, srcId)]
+    True -> case gtnParent node of
+      Nothing -> [(srcId, srcId)]
+      Just parentId -> [(srcId, parentId)]
+  _ -> []  -- not a lexeme
+
+renderNamingInfo :: NamingInfo -> RenderedOutputs
+renderNamingInfo _ = mempty
+
+renderTypeInfo :: TypeInfo -> RenderedOutputs
+renderTypeInfo _ = mempty
+
+
 instance ToJSON RenderedSourceBlock
 instance ToJSON RenderedOutput
+instance ToJSON TreeNodeState
+instance ToJSON TreeNodeUpdate
+instance ToJSON HighlightType
 
 -- -----------------
 
diff --git a/src/lib/Runtime.hs b/src/lib/Runtime.hs
index 2d0a900f..c9a33638 100644
--- a/src/lib/Runtime.hs
+++ b/src/lib/Runtime.hs
@@ -77,7 +77,7 @@ checkedCallFunPtr fd argsPtr resultPtr fPtr = do
 withPipeToLogger :: PassLogger -> (FD -> IO a) -> IO a
 withPipeToLogger logger writeAction = do
   result <- snd <$> withPipe
-    (\h -> readStream h \s -> ioLogAction logger [TextOut s])
+    (\h -> readStream h \s -> ioLogAction logger (Outputs [TextOut s]))
     (\h -> handleToFd h >>= writeAction)
   case result of
     Left e -> E.throw e
diff --git a/src/lib/SourceIdTraversal.hs b/src/lib/SourceIdTraversal.hs
index 7e243620..8333bc78 100644
--- a/src/lib/SourceIdTraversal.hs
+++ b/src/lib/SourceIdTraversal.hs
@@ -4,15 +4,27 @@
 -- license that can be found in the LICENSE file or at
 -- https://developers.google.com/open-source/licenses/bsd
 
-module SourceIdTraversal (getGroupTree) where
+module SourceIdTraversal (getGroupingInfo) where
 
 import Control.Monad.Writer.Strict
+import qualified Data.Map.Strict as M
 import Data.Functor ((<&>))
 
 import Types.Source
 import Types.Primitives
 import Err
 
+getGroupingInfo :: SourceBlock' -> GroupingInfo
+getGroupingInfo sb = groupTreeToGroupingInfo $ getGroupTree sb
+
+groupTreeToGroupingInfo :: GroupTree -> GroupingInfo
+groupTreeToGroupingInfo groupTreeTop = execWriter $ go Nothing groupTreeTop where
+  go :: Maybe SrcId -> GroupTree -> Writer GroupingInfo ()
+  go parent (GroupTree sid lexSpan children isAtomic) = do
+    mapM_ (go (Just sid)) children
+    let node = GroupTreeNode parent lexSpan (map gtSrcId children) isAtomic
+    tell $ GroupingInfo $ M.singleton sid node
+
 getGroupTree :: SourceBlock' -> GroupTree
 getGroupTree b = mkGroupTree False rootSrcId $ runTreeM $ visit b
 
diff --git a/src/lib/SourceRename.hs b/src/lib/SourceRename.hs
index d5420dab..aadeca08 100644
--- a/src/lib/SourceRename.hs
+++ b/src/lib/SourceRename.hs
@@ -25,7 +25,7 @@ import Types.Primitives
 import Types.Top (Env (..), ModuleEnv (..))
 
 renameSourceNamesTopUDecl
-  :: (Fallible1 m, EnvReader m)
+  :: (Fallible1 m, EnvReader m, TopLogger1 m)
   => ModuleSourceName -> UTopDecl VoidS VoidS -> m n (Abs UTopDecl SourceMap n)
 renameSourceNamesTopUDecl mname decl = do
   Distinct <- getDistinct
@@ -40,7 +40,7 @@ uDeclErrSourceMap mname decl =
   SourceMap $ M.fromSet (const [ModuleVar mname Nothing]) (sourceNames decl)
 {-# SCC uDeclErrSourceMap #-}
 
-renameSourceNamesUExpr :: (Fallible1 m, EnvReader m) => UExpr VoidS -> m n (UExpr n)
+renameSourceNamesUExpr :: (Fallible1 m, EnvReader m, TopLogger1 m) => UExpr VoidS -> m n (UExpr n)
 renameSourceNamesUExpr expr = do
   Distinct <- getDistinct
   liftRenamer $ sourceRenameE expr
diff --git a/src/lib/TopLevel.hs b/src/lib/TopLevel.hs
index 80ae85d7..8148ff03 100644
--- a/src/lib/TopLevel.hs
+++ b/src/lib/TopLevel.hs
@@ -54,7 +54,6 @@ import Err
 import IRVariants
 import Imp
 import ImpToLLVM
-import IncState
 import Inference
 import Inline
 import Lower
@@ -219,7 +218,7 @@ evalSourceBlock
   :: (Topper m, Mut n) => ModuleSourceName -> SourceBlock -> m n (Except ())
 evalSourceBlock mname block = do
   maybeErr <- catchErrExcept do
-    logTop $ SourceInfo $ SIGroupTree $ OverwriteWith $ getGroupTree $ sbContents block
+    logTop $ SourceInfo $ SIGroupingInfo $ getGroupingInfo $ sbContents block
     evalSourceBlock' mname block
   case (maybeErr, sbContents block) of
     (Failure _, TopDecl decl) -> do
@@ -669,14 +668,14 @@ checkPass name cont = do
   return result
 
 logTop :: TopLogger m => Output -> m ()
-logTop x = emitLog [x]
+logTop x = emitLog $ Outputs [x]
 
 logDebug :: TopLogger m => m Output -> m ()
 logDebug m = getLogLevel >>= \case
   NormalLogLevel -> return ()
   DebugLogLevel -> do
     x <- m
-    emitLog [x]
+    emitLog $ Outputs [x]
 
 logPass :: Topper m => Pretty a => PassName -> a -> m n ()
 logPass passName result = do
diff --git a/src/lib/Types/Source.hs b/src/lib/Types/Source.hs
index 0cae1dc9..83b37531 100644
--- a/src/lib/Types/Source.hs
+++ b/src/lib/Types/Source.hs
@@ -42,7 +42,6 @@ import qualified Types.OpNames as P
 import IRVariants
 import MonadUtil
 import Util (File (..), SnocList)
-import IncState
 
 import Types.Primitives
 
@@ -97,19 +96,50 @@ instance Monoid LexemeInfo where
 
 -- === Source info ===
 
-newtype TypeInfo = TypeInfo { fromTypeInfo :: M.Map SrcId String }
-        deriving (Semigroup, Monoid, Show, Eq)
-
 data SourceInfo =
-   SIGroupTree (Overwrite GroupTree)
- | SITypeInfo  TypeInfo
-     deriving (Show, Eq, Generic)
+   SIGroupingInfo  GroupingInfo
+ | SINamingInfo    NamingInfo
+ | SITypeInfo      TypeInfo
+   deriving (Show, Eq, Generic)
+
+newtype GroupingInfo = GroupingInfo (M.Map SrcId GroupTreeNode)
+        deriving (Show, Eq, Semigroup, Monoid, Generic)
+data GroupTreeNode = GroupTreeNode
+  { gtnParent         :: Maybe SrcId
+  , gtnSpan           :: LexemeSpan
+  , gtnChildren       :: [SrcId]
+  , gtnIsAtomicLexeme :: Bool }
+    deriving (Show, Eq, Generic)
+
+data NamingInfo = NamingInfo (M.Map SrcId NameInfo)
+  deriving (Show, Eq, Generic)
+data NameInfo =
+   LocalBinder [SrcId] -- src ids of groups for which this binder is in scope
+ | LocalOcc SrcId      -- src id of this occ's binder
+ | TopOcc String
+ deriving (Show, Eq, Generic)
+
+newtype TypeInfo = TypeInfo (M.Map SrcId TypeInfo)
+        deriving (Show, Eq, Semigroup, Monoid, Generic)
+type TypeStr = String
+type ExprStr = String
+data NodeTypeInfo =
+   ExprType TypeStr      -- type of arbitrary expression
+ | BinderType TypeStr
+ | AppType
+    TypeStr             -- type of whole application expression
+    [(String, TypeStr)] -- names and inferred types of implicit args
+    [ExprStr]           -- values of synthesized dictionaries
+    [SrcId]             -- binder srcIds for vars ocurring in terms produce by inference
+  deriving (Show, Eq, Generic)
 
 -- === Results ===
 
+type TopLogger1 (m::MonadKind1) = forall n. Logger Outputs (m n)
+
 type LitProg = [(SourceBlock, Outputs)]
 
-type Outputs = [Output]
+newtype Outputs = Outputs [Output] deriving (Show, Eq, Generic, Semigroup, Monoid)
 data Output =
     TextOut String
   | HtmlOut String
@@ -943,13 +973,8 @@ deriving instance Eq   (UEffectRow n)
 deriving instance Ord  (UEffectRow n)
 
 instance ToJSON LexemeType
-instance ToJSON SourceInfo
-instance ToJSON GroupTree
 instance ToJSON PassName
 
-instance ToJSON TypeInfo where
-  toJSON m = toJSON $ M.toList $ fromTypeInfo m
-
 -- === Pretty instances ===
 
 instance Pretty CSBlock where
diff --git a/static/index.ts b/static/index.ts
index 176bb8cd..6b0320a0 100644
--- a/static/index.ts
+++ b/static/index.ts
@@ -13,15 +13,24 @@ const minimap      : Element = document.getElementById("minimap") ?? oops()
 
 type CellId = number
 type SrcId  = number
+type LexemeId  = number
 type HTMLString = string
 type Div = Element
 type Status = "Waiting" | "Running" | "Complete" | "CompleteWithErrors" | "Inert"
-type HighlightType = "HighlightGroup" |  "HighlightLeaf"
-type Span = [Div, Div]
-type SpanMap      = Map<SrcId, Span>
-type HighlightMap = Map<SrcId, [HighlightType, SrcId]>
-type HoverMap     = Map<SrcId, HTMLString>
 
+type HighlightType = "HighlightGroup" |  "HighlightLeaf" | "HighlightError"
+type Highlight    = [SrcId, HighlightType]
+
+type HsMaybe<T> = {tag:"Just"; contents:T} | {tag: "Nothing"}
+type HsOverwrite<T> = {tag:"OverwriteWith"; contents:T} | {tag: "NoChange"}
+
+type FocusMap = Map<LexemeId, SrcId>
+type TreeMap  = Map<SrcId, TreeNode>
+type TreeNode = {
+    srcId      : SrcId
+    span       : [Div, Div] | null ;
+    highlights : Highlight[];
+    text       : HTMLString}
 type Cell = {
     cellId       : CellId;
     root         : Div;
@@ -31,12 +40,8 @@ type Cell = {
     status       : Div;
     curStatus    : Status | null ;
     sourceText   : string;
-    spanMap      : SpanMap;
-    highlightMap : HighlightMap;
-    hoverMap     : HoverMap}
-type HsSourceInfo =
-    {tag: "SIGroupTree"; contents: HsMaybe<HsGroupTree>} |
-    {tag: "SITypeInfo" ; contents: [SrcId, string][]}
+    focusMap     : FocusMap;
+    treeMap      : TreeMap}
 type HsRenderedSourceBlock = {
     rsbLine       : number;
     rsbNumLines   : number;
@@ -47,22 +52,25 @@ type HsRenderedSourceBlock = {
 type HsRenderedOutput =
     {tag: "RenderedTextOut"   ; contents: string                  } |
     {tag: "RenderedHtmlOut"   ; contents: HTMLString              } |
-    {tag: "RenderedSourceInfo"; contents: HsSourceInfo            } |
-    {tag: "RenderedPassResult"; contents: PassName                } |
+    {tag: "RenderedPassResult"; contents: string                  } |
     {tag: "RenderedMiscLog"   ; contents: string                  } |
-    {tag: "RenderedError"     ; contents: [HsMaybe<SrcId>, string]}
-type PassName = string
-type PassInfo = string | null
-type LexemeId = SrcId
-type LexemeSpan = [LexemeId, LexemeId]
-type HsGroupTree = {
-    gtSrcId          : SrcId;
-    gtSpan           : LexemeSpan;
-    gtChildren       : [HsGroupTree];
-    gtIsAtomicLexeme : boolean }
+    {tag: "RenderedError"     ; contents: [HsMaybe<SrcId>, string]} |
+    {tag: "RenderedTreeNodeUpdate" ; contents: [SrcId, HsTreeNodeMapUpdate][]} |
+    {tag: "RenderedFocusUpdate"    ; contents: [LexemeId, SrcId][]}
 
-type HsMaybe<T> = {tag:"Just"; contents:T} | {tag: "Nothing"}
-type HsOverwrite<T> = {tag:"OverwriteWith"; contents:T} | {tag: "NoChange"}
+type HsFocusMap = Map<LexemeId, SrcId>
+type HsTreeNodeState = {
+    tnSpan        : [LexemeId, LexemeId]
+    tnHighlights  : Highlight[]
+    tnText        : HTMLString}
+type HsTreeNodeUpdate = {
+    tnuHighlights : HsOverwrite<Highlight[]>;
+    tnuText       : HsOverwrite<HTMLString>}
+type HsTreeNodeMapUpdate =
+    {tag: "Create" ; contents: HsTreeNodeState}  |
+    {tag: "Replace"; contents: HsTreeNodeState}  |
+    {tag: "Update" ; contents: HsTreeNodeUpdate} |
+    {tag: "Delete"}
 type HsCellState = [HsRenderedSourceBlock, Status, HsRenderedOutput[]]
 type HsCellUpdate = [HsOverwrite<Status>, HsRenderedOutput[]]
 type HsCellMapUpdate =
@@ -147,9 +155,8 @@ function createCell(cellId: CellId) : Cell {
         status       : status,
         curStatus    : null,
         sourceText   : "",
-        spanMap      : new Map<SrcId, Span>(),
-        highlightMap : new Map<SrcId, [HighlightType, SrcId]>(),
-        hoverMap     : new Map<SrcId, HTMLString>()}
+        focusMap     : new Map<LexemeId, SrcId>(),
+        treeMap      : new Map<SrcId, TreeNode>()}
     cells.set(cellId, cell)
     return cell
 }
@@ -173,7 +180,7 @@ function updateCellContents(cell:Cell, update:HsCellUpdate) {
 function removeHover() {
     hoverInfoDiv.innerHTML = ""
     curHighlights.map(function (element) {
-        element.classList.remove("highlighted", "highlighted-leaf")})
+        element.classList.remove("highlight-group", "highlight-leaf")})
     curHighlights.length = 0
 }
 function attachStatusHovertip(cell:Cell) {
@@ -186,12 +193,13 @@ function attachStatusHovertip(cell:Cell) {
 }
 function attachHovertip(cell:Cell, srcId:SrcId) {
     let span = selectSpan(cell, srcId)
-    span.addEventListener("mouseover", function (event:Event) {
-        event.stopPropagation()
-        applyCellHover(cell, srcId)})
-    span.addEventListener("mouseout" , function (event:Event) {
-        event.stopPropagation()
-        removeHover()})
+    if (span !== null) {
+        span.addEventListener("mouseover", function (event:Event) {
+            event.stopPropagation()
+            applyCellHover(cell, srcId)})
+        span.addEventListener("mouseout" , function (event:Event) {
+            event.stopPropagation()
+            removeHover()})}
 }
 function addChild(div:Div, className:string) : Div {
     const child = document.createElement("div")
@@ -202,18 +210,9 @@ function addChild(div:Div, className:string) : Div {
 function appendText(div:Div, s:string) {
     div.appendChild(document.createTextNode(s))
 }
-function selectSpan(cell:Cell, srcId:SrcId) : Div {
+function selectSpan(cell:Cell, srcId:SrcId) : Div | null {
     const spanId : string = "#span_".concat(cell.cellId.toString(), "_", srcId.toString())
-    return cell.source.querySelector(spanId) ?? oops()
-}
-function addClassToSrcNode(cell: Cell, srcId:SrcId, className:string) {
-    // let span = getSpan(cellId, srcId)
-    // if (span !== undefined) {
-    //     let [l, r] = span
-    //     let spans = spansBetween(selectSpan(cellId, l), selectSpan(cellId, r));
-    //     spans.map(function (span) {
-    //         span.classList.add(className)
-    //         curHighlights.push(span)})}
+    return cell.source.querySelector(spanId)
 }
 function spansBetween(l:Div, r:Div) : Div[] {
     let spans : Div[] = []
@@ -269,9 +268,6 @@ function extendCellOutput(cell: Cell, outputs:HsRenderedOutput[]) {
             case "RenderedHtmlOut":
                 addHTMLResult(cell, output.contents)
                 break
-            case "RenderedSourceInfo":
-                extendSourceInfo(cell, output.contents)
-                break
             case "RenderedPassResult":
                 // TODO: show passes
                 break
@@ -281,37 +277,93 @@ function extendCellOutput(cell: Cell, outputs:HsRenderedOutput[]) {
             case "RenderedError":
                 const [maybeSrcId, errMsg] = output.contents
                 if (maybeSrcId.tag == "Just") {
-                    addClassToSrcNode(cell, maybeSrcId.contents, "err-span")}
+                    const node : TreeNode = cell.treeMap.get(maybeSrcId.contents) ?? oops()
+                    highlightTreeNode(false, node, "HighlightError")}
                 addErrResult(cell, errMsg)
                 break
+            case "RenderedTreeNodeUpdate":
+                output.contents.forEach(function (elt:[SrcId, HsTreeNodeMapUpdate]) {
+                    const [srcId, update] = elt
+                    applyTreeNodeUpdate(cell, srcId, update)})
+                break
+            case "RenderedFocusUpdate":
+                output.contents.forEach(function (elt:[LexemeId, SrcId]) {
+                    const [lexemeId, srcId] = elt
+                    cell.focusMap.set(lexemeId, srcId)})
+                break
             default:
                 // nothing
         }})
 }
-function extendSourceInfo(cell: Cell, info: HsSourceInfo) {
-    switch (info.tag) {
-        case "SITypeInfo":
-            // TODO: this should really merge with previous rather than
-            // clobbering it completely but for now we only expect to do this
-            // once per cell so it's ok.
-            cell.hoverMap = new Map(info.contents)
+function applyTreeNodeUpdate(cell:Cell, srcId:SrcId, update:HsTreeNodeMapUpdate) {
+    switch (update.tag) {
+        case "Create":  // deliberate fallthrough
+        case "Replace":
+            const s : HsTreeNodeState = update.contents
+            const [l, r] = s.tnSpan
+            const range = computeRange(cell, l, r)
+            const treeNode : TreeNode = {
+                srcId      : srcId,
+                span       : range,
+                highlights : s.tnHighlights,
+                text       : s.tnText}
+            cell.treeMap.set(srcId, treeNode)
             break
-        default:
-            // nothing
+        case "Update":
+            const nodeUpdate : HsTreeNodeUpdate = update.contents
+            const curTreeNode : TreeNode = cell.treeMap.get(srcId) ?? oops()
+            const text = nodeUpdate.tnuText
+            if (text.tag == "OverwriteWith") {
+                curTreeNode.text = text.contents}
+            const hl = nodeUpdate.tnuHighlights
+            if (hl.tag == "OverwriteWith") {
+                curTreeNode.highlights = hl.contents}
+            break}
+}
+function computeRange(cell:Cell, l:SrcId, r:SrcId) : [Div, Div] | null {
+    const lDiv = selectSpan(cell, l)
+    const rDiv = selectSpan(cell, r)
+    if (lDiv !== null && rDiv !== null) {
+        return [lDiv, rDiv]
+    } else {
+        return null}
+}
+function applyCellHover(cell:Cell, srcId:LexemeId) {
+    const focus : undefined | SrcId = cell.focusMap.get(srcId)
+    if (focus !== undefined) {
+        applyFocus(cell, focus)
     }
 }
-function applyCellHover(cell:Cell, srcId:SrcId) {
-    applyHoverHighlights(cell, srcId)
-    const hoverInfo : undefined | HTMLString = cell.hoverMap.get(srcId)
-    if (hoverInfo !== undefined) {
-        setHoverInfo(hoverInfo)}
+function applyFocus(cell:Cell, focus:SrcId) {
+    const focusNode : TreeNode = cell.treeMap.get(focus) ?? oops()
+    focusNode.highlights.forEach((h:Highlight) => {
+        const [sid, ty] = h
+        const node : TreeNode = cell.treeMap.get(sid) ?? oops()
+        highlightTreeNode(true, node, ty)})
+    setHoverInfo(focusNode.text)
 }
 function setHoverInfo(s:string) {
     hoverInfoDiv.innerHTML = ""
     appendText(hoverInfoDiv, s)
 }
-function applyHoverHighlights(cell:Cell, srcId:SrcId) {
-    // TODO
+function computeHighlightClass(h:HighlightType) : string {
+    switch (h) {
+        case "HighlightGroup":
+            return "highlight-group"
+        case "HighlightLeaf":
+            return "highlight-leaf"
+        case "HighlightError":
+            return "highlight-error"
+    }
+}
+function highlightTreeNode(isTemporary: boolean, node: TreeNode, highlightType:HighlightType) {
+    const highlightClass : string = computeHighlightClass(highlightType)
+    if (node.span !== null) {
+        let [l, r] = node.span
+        let spans = spansBetween(l, r)
+        spans.map(function (span) {
+            span.classList.add(highlightClass)
+            if (isTemporary) {curHighlights.push(span)}})}
 }
 type RenderMode = "Static" | "Dynamic"
 function render(renderMode:RenderMode) {
diff --git a/static/style.css b/static/style.css
index cafc2f5d..a2860b23 100644
--- a/static/style.css
+++ b/static/style.css
@@ -43,25 +43,25 @@ body {
   font-family: monospace;
   white-space: pre;
 }
-
 /* cell structure */
 .cell {
   margin-left: 5px;
   display: flex;
 }
-.line-nums, .contents {
-  font-family: monospace;
-  white-space: pre;
-}
 .line-nums {
   flex: 0 0 3em;
   height: 100%;
   text-align: right;
   color: #808080;
+  font-family: monospace;
+  white-space: pre;
 }
 .contents {
   margin-left: 1em;
+  font-family: monospace;
+  white-space: pre;
 }
+
 /* special results */
 .err-result {
   font-weight: bold;
@@ -69,19 +69,20 @@ body {
 }
 
 /* status colors */
+.status-inert   {}
 .status-waiting {background-color: gray;}
 .status-running {background-color: lightblue;}
 .status-err     {background-color: red;}
-.status-inert   {}
 .status-success {background-color: white;}
 
-/* error highlighting */
-.err-span {
+/* span highlighting */
+.highlight-error {
   text-decoration: red wavy underline;
-  text-decoration-skip-ink: none;
-}
+  text-decoration-skip-ink: none;}
+.highlight-group { background-color: yellow; }
+.highlight-leaf  { background-color: lightgray; }
 
-/* Lexeme colors */
+/* lexeme colors */
 .comment   {color: gray;}
 .keyword   {color: #0000DD;}
 .command   {color: #A80000;}