@@ -72,7 +72,7 @@ func FindRangesFromIndexList(stack *nodestack.NodeStack, indexList []string, vm
72
72
return FindRangesFromIndexList (stack , indexList , vm )
73
73
case * ast.Function :
74
74
// If the function's body is an object, it means we can look for indexes within the function
75
- if funcBody , ok := bodyNode .Body .( * ast. DesugaredObject ); ok {
75
+ if funcBody := findChildDesugaredObject ( bodyNode .Body ); funcBody != nil {
76
76
foundDesugaredObjects = append (foundDesugaredObjects , funcBody )
77
77
}
78
78
default :
@@ -124,34 +124,29 @@ func extractObjectRangesFromDesugaredObjs(stack *nodestack.NodeStack, vm *jsonne
124
124
return nil , err
125
125
}
126
126
// If the reference is an object, add it directly to the list of objects to look in
127
- if varReference , ok := varReference .( * ast. DesugaredObject ); ok {
127
+ if varReference := findChildDesugaredObject ( varReference ); varReference != nil {
128
128
desugaredObjs = append (desugaredObjs , varReference )
129
129
}
130
130
// If the reference is a function, and the body of that function is an object, add it to the list of objects to look in
131
131
if varReference , ok := varReference .(* ast.Function ); ok {
132
- if funcBody , ok := varReference .Body .( * ast. DesugaredObject ); ok {
132
+ if funcBody := findChildDesugaredObject ( varReference .Body ); funcBody != nil {
133
133
desugaredObjs = append (desugaredObjs , funcBody )
134
134
}
135
135
}
136
136
case * ast.DesugaredObject :
137
- stack .Push (fieldNode )
138
- desugaredObjs = append (desugaredObjs , findDesugaredObjectFromStack (stack ))
137
+ desugaredObjs = append (desugaredObjs , fieldNode )
139
138
case * ast.Index :
140
- tempStack := nodestack .NewNodeStack (fieldNode )
141
- additionalIndexList := tempStack .BuildIndexList ()
142
- additionalIndexList = append (additionalIndexList , indexList ... )
139
+ additionalIndexList := append (nodestack .NewNodeStack (fieldNode ).BuildIndexList (), indexList ... )
143
140
result , err := FindRangesFromIndexList (stack , additionalIndexList , vm )
144
- if sameFileOnly && len (result ) > 0 && result [0 ].Filename != stack .From .Loc ().FileName {
145
- continue
146
- }
147
141
if len (result ) > 0 {
148
- return result , err
142
+ if ! sameFileOnly || result [0 ].Filename == stack .From .Loc ().FileName {
143
+ return result , err
144
+ }
149
145
}
150
146
151
147
fieldNodes = append (fieldNodes , fieldNode .Target )
152
148
case * ast.Function :
153
- stack .Push (fieldNode .Body )
154
- desugaredObjs = append (desugaredObjs , findDesugaredObjectFromStack (stack ))
149
+ desugaredObjs = append (desugaredObjs , findChildDesugaredObject (fieldNode .Body ))
155
150
case * ast.Import :
156
151
filename := fieldNode .File .Value
157
152
newObjs := findTopLevelObjectsInFile (vm , filename , string (fieldNode .Loc ().File .DiagnosticFileName ))
@@ -222,14 +217,16 @@ func findObjectFieldInObject(objectNode *ast.DesugaredObject, index string) *ast
222
217
return nil
223
218
}
224
219
225
- func findDesugaredObjectFromStack (stack * nodestack.NodeStack ) * ast.DesugaredObject {
226
- for ! stack .IsEmpty () {
227
- switch node := stack .Pop ().(type ) {
228
- case * ast.DesugaredObject :
229
- return node
230
- case * ast.Binary :
231
- stack .Push (node .Left )
232
- stack .Push (node .Right )
220
+ func findChildDesugaredObject (node ast.Node ) * ast.DesugaredObject {
221
+ switch node := node .(type ) {
222
+ case * ast.DesugaredObject :
223
+ return node
224
+ case * ast.Binary :
225
+ if res := findChildDesugaredObject (node .Left ); res != nil {
226
+ return res
227
+ }
228
+ if res := findChildDesugaredObject (node .Right ); res != nil {
229
+ return res
233
230
}
234
231
}
235
232
return nil
@@ -262,7 +259,7 @@ func findLHSDesugaredObject(stack *nodestack.NodeStack) (*ast.DesugaredObject, e
262
259
case * ast.Var :
263
260
bind := FindBindByIDViaStack (stack , lhsNode .Id )
264
261
if bind != nil {
265
- if bindBody , ok := bind .Body .( * ast. DesugaredObject ); ok {
262
+ if bindBody := findChildDesugaredObject ( bind .Body ); bindBody != nil {
266
263
return bindBody , nil
267
264
}
268
265
}
0 commit comments