content: -- Copyright 2023 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  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
  
  type TreeM = Writer [GroupTree]
  
  mkGroupTree :: Bool -> SrcId -> [GroupTree] -> GroupTree
  mkGroupTree isAtomic sid = \case
    [] -> GroupTree sid (sid,sid) [] isAtomic -- no children - must be a lexeme
    subtrees -> GroupTree sid (l,r) subtrees isAtomic
      where l = minimum $ subtrees <&> (fst . gtSpan)
            r = maximum $ subtrees <&> (snd . gtSpan)
  
  runTreeM :: TreeM () -> [GroupTree]
  runTreeM cont = snd $ runWriter $ cont
  
  enterNode :: SrcId -> TreeM () -> TreeM ()
  enterNode sid cont = tell [mkGroupTree False sid (runTreeM cont)]
  
  emitLexeme :: SrcId -> TreeM ()
  emitLexeme lexemeId = tell [mkGroupTree True lexemeId []]
  
  class IsTree a where
    visit :: a -> TreeM ()
  
  instance IsTree SourceBlock' where
    visit = \case
      TopDecl decl -> visit decl
      Command _ g -> visit g
      DeclareForeign v1 v2 g -> visit v1 >> visit v2 >> visit g
      DeclareCustomLinearization v _ g -> visit v >> visit g
      Misc _ -> return ()
      UnParseable _ _ -> return ()
  
  instance IsTree Group where
    visit = \case
      CLeaf _ -> return ()
      CPrim _   xs -> mapM_ visit xs
      CParens   xs -> mapM_ visit xs
      CBrackets xs -> mapM_ visit xs
      CBin       b l r -> visit l >> visit b >> visit r
      CJuxtapose _ l r -> visit l >> visit r
      CPrefix      l r -> visit l >> visit r
      CGivens (x,y) -> visit x >> visit y
      CLambda args body -> visit args >> visit body
      CFor _ args body -> visit args >> visit body
      CCase scrut alts -> visit scrut >> visit alts
      CIf scrut ifTrue ifFalse -> visit scrut >> visit ifTrue >> visit ifFalse
      CDo body -> visit body
      CArrow l effs r -> visit l >> visit effs >> visit r
      CWith b body -> visit b >> visit body
  
  instance IsTree Bin where
    visit = \case
      EvalBinOp b -> visit b
      _ -> return ()
  
  instance IsTree CSBlock where
    visit = \case
      IndentedBlock sid decls -> enterNode sid $ visit decls
      ExprBlock body -> visit body
  
  instance IsTree CSDecl where
    visit = \case
      CLet v rhs -> visit v >> visit rhs
      CDefDecl def -> visit def
      CExpr g -> visit g
      CBind v body -> visit v >> visit body
      CPass -> return ()
  
  instance IsTree CTopDecl where
    visit = \case
      CSDecl _ decl -> visit decl
      CData v params givens cons -> visit v >> visit params >> visit givens >> visit cons
      CStruct v params givens fields methods -> visit v >> visit params >> visit givens >> visit fields >> visit methods
      CInterface v params methods -> visit v >> visit params >> visit methods
      CInstanceDecl def -> visit def
  
  instance IsTree CDef where
    visit (CDef v params rhs givens body) =
      visit v >> visit params >> visit rhs >> visit givens >> visit body
  
  instance IsTree CInstanceDef where
    visit (CInstanceDef v args givens methods name) =
      visit v >> visit args >> visit givens >> visit methods >> visit name
  
  instance IsTree a => IsTree (WithSrc a) where
    visit (WithSrc sid x) = enterNode sid $ visit x
  
  instance IsTree a => IsTree (WithSrcs a) where
    visit (WithSrcs sid sids x) = enterNode sid $ mapM_ emitLexeme sids >> visit x
  
  instance IsTree a => IsTree [a] where
    visit xs = mapM_ visit xs
  
  instance IsTree a => IsTree (Maybe a) where
    visit xs = mapM_ visit xs
  
  instance (IsTree a, IsTree b) => IsTree (a, b) where
    visit (x, y) = visit x >> visit y
  
  instance (IsTree a, IsTree b, IsTree c) => IsTree (a, b, c) where
    visit (x, y, z) = visit x >> visit y >> visit z
  
  instance IsTree AppExplicitness where visit _ = return ()
  instance IsTree SourceName      where visit _ = return ()
  instance IsTree LetAnn          where visit _ = return ()
  