From ab95594eea6510c0678cffc45df2c6cf7b853ec4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stanis=C5=82aw=20Barzowski?= <stanislaw.barzowski@gmail.com>
Date: Mon, 9 Sep 2019 17:25:12 +0200
Subject: [PATCH] Add missing cases in specialChildren

Namely:
1) function body and optional params
2) object locals
---
 internal/parser/context.go | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/internal/parser/context.go b/internal/parser/context.go
index 07984ce..642ac6b 100644
--- a/internal/parser/context.go
+++ b/internal/parser/context.go
@@ -98,7 +98,7 @@ func directChildren(node ast.Node) []ast.Node {
 	case *ast.Object:
 		return objectFieldsDirectChildren(node.Fields)
 	case *ast.DesugaredObject:
-		return desugaredObjectFieldsDirectChildren(node.Fields)
+		return desugaredObjectDirectChildren(node)
 	case *ast.ArrayComp:
 		result := []ast.Node{}
 		spec := &node.Spec
@@ -247,23 +247,29 @@ func inObjectFieldsChildren(fields ast.ObjectFields) ast.Nodes {
 	return result
 }
 
-func desugaredObjectFieldsDirectChildren(fields ast.DesugaredObjectFields) ast.Nodes {
+func desugaredObjectDirectChildren(obj *ast.DesugaredObject) ast.Nodes {
 	result := ast.Nodes{}
-	for _, field := range fields {
+	for _, field := range obj.Fields {
+		if field.Name == nil {
+			panic("Name cannot be nil")
+		}
 		result = append(result, field.Name)
 	}
 	return result
 }
 
-func inDesugaredObjectFieldsChildren(fields ast.DesugaredObjectFields) ast.Nodes {
-	result := ast.Nodes{}
-	for _, field := range fields {
+func inDesugaredObjectSpecialChildren(obj *ast.DesugaredObject) ast.Nodes {
+	result := make([]ast.Node, 0, len(obj.Fields)+len(obj.Locals))
+	for _, field := range obj.Fields {
 		result = append(result, field.Body)
 	}
+	for _, local := range obj.Locals {
+		result = append(result, local.Body)
+	}
 	return result
 }
 
-// specialChildren returns children are neither direct nor thunked,
+// specialChildren returns children that are neither direct nor thunked,
 // e.g. object field body.
 // These nodes are evaluated in a different environment from their parent.
 //
@@ -287,8 +293,11 @@ func specialChildren(node ast.Node) []ast.Node {
 	case *ast.Error:
 		return nil
 	case *ast.Function:
-		// TODO(sbarzowski) this
-		return nil
+		children := []ast.Node{node.Body}
+		for _, child := range node.Parameters.Optional {
+			children = append(children, child.DefaultArg)
+		}
+		return children
 	case *ast.Import:
 		return nil
 	case *ast.ImportStr:
@@ -313,7 +322,7 @@ func specialChildren(node ast.Node) []ast.Node {
 	case *ast.LiteralString:
 		return nil
 	case *ast.DesugaredObject:
-		return inDesugaredObjectFieldsChildren(node.Fields)
+		return inDesugaredObjectSpecialChildren(node)
 	case *ast.Object:
 		return inObjectFieldsChildren(node.Fields)
 	case *ast.ArrayComp: