summaryrefslogtreecommitdiff
path: root/src/lib/SourceIdTraversal.hs
blob: 9a1a9e09b0da975a75e6a5ffdfaad5e8c175462a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
-- 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 ()