Added decorating and restoring Objects and Scopes

decorator/decorator.go

@@ -1,6 +1,7 @@
 package decorator
 
 import (
+	"fmt"
 	"go/ast"
 	"go/parser"
 	"go/token"
@@ -66,12 +67,16 @@ func DecorateFile(fset *token.FileSet, f *ast.File) *dst.File {
 func New() *Decorator {
 	return &Decorator{
 		Nodes:       map[ast.Node]dst.Node{},
+		Scopes:      map[*ast.Scope]*dst.Scope{},
+		Objects:     map[*ast.Object]*dst.Object{},
 		decorations: map[ast.Node]map[string][]string{},
 	}
 }
 
 type Decorator struct {
 	Nodes       map[ast.Node]dst.Node
+	Scopes      map[*ast.Scope]*dst.Scope
+	Objects     map[*ast.Object]*dst.Object
 	decorations map[ast.Node]map[string][]string
 }
 
@@ -86,3 +91,89 @@ type decorationInfo struct {
 	name string
 	decs []string
 }
+
+func (d *Decorator) decorateObject(o *ast.Object) *dst.Object {
+	if o == nil {
+		return nil
+	}
+	if do, ok := d.Objects[o]; ok {
+		return do
+	}
+	/*
+		// An Object describes a named language entity such as a package,
+		// constant, type, variable, function (incl. methods), or label.
+		//
+		// The Data fields contains object-specific data:
+		//
+		//	Kind    Data type         Data value
+		//	Pkg     *Scope            package scope
+		//	Con     int               iota for the respective declaration
+		//
+		type Object struct {
+			Kind ObjKind
+			Name string      // declared name
+			Decl interface{} // corresponding Field, XxxSpec, FuncDecl, LabeledStmt, AssignStmt, Scope; or nil
+			Data interface{} // object-specific data; or nil
+			Type interface{} // placeholder for type information; may be nil
+		}
+	*/
+
+	out := &dst.Object{}
+	d.Objects[o] = out
+	out.Kind = dst.ObjKind(o.Kind)
+	out.Name = o.Name
+
+	switch decl := o.Decl.(type) {
+	case *ast.Scope:
+		out.Decl = d.decorateScope(decl)
+	case ast.Node:
+		out.Decl = d.decorateNode(decl)
+	default:
+		panic(fmt.Sprintf("o.Decl is %T", o.Decl))
+	}
+
+	// TODO: I believe Data is either a *Scope or an int. We will support both and panic if something else if found.
+	switch data := o.Data.(type) {
+	case int:
+		out.Data = data
+	case *ast.Scope:
+		out.Data = d.decorateScope(data)
+	case ast.Node:
+		out.Data = d.decorateNode(data)
+	case nil:
+	default:
+		panic(fmt.Sprintf("o.Data is %T", o.Data))
+	}
+
+	return out
+}
+
+func (d *Decorator) decorateScope(s *ast.Scope) *dst.Scope {
+	if s == nil {
+		return nil
+	}
+	if ds, ok := d.Scopes[s]; ok {
+		return ds
+	}
+	/*
+		// A Scope maintains the set of named language entities declared
+		// in the scope and a link to the immediately surrounding (outer)
+		// scope.
+		//
+		type Scope struct {
+			Outer   *Scope
+			Objects map[string]*Object
+		}
+	*/
+	out := &dst.Scope{}
+
+	d.Scopes[s] = out
+
+	out.Outer = d.decorateScope(s.Outer)
+	out.Objects = map[string]*dst.Object{}
+	for k, v := range s.Objects {
+		out.Objects[k] = d.decorateObject(v)
+	}
+
+	return out
+}

decorator/decorator_test.go

@@ -20,6 +20,22 @@ func TestDecorator(t *testing.T) {
 		code       string
 		expect     string
 	}{
+		{
+			name: "range bug",
+			code: `package main
+
+				func main() {
+					/*Start*/
+					for /*AfterFor*/ k /*AfterKey*/, v /*AfterValue*/ := range /*AfterRange*/ a /*AfterX*/ {
+						print(k, v)
+					} /*End*/
+				}`,
+			expect: `File [AfterName "\n" "\n"]
+            BlockStmt [AfterLbrace "\n"]
+            RangeStmt [Start "/*Start*/" "\n"] [AfterFor "/*AfterFor*/"] [AfterKey "/*AfterKey*/"] [AfterValue "/*AfterValue*/"] [AfterRange "/*AfterRange*/"] [AfterX "/*AfterX*/"] [End "/*End*/" "\n"]
+            BlockStmt [AfterLbrace "\n"]
+            ExprStmt [End "\n"]`,
+		},
 		{
 			name: "value spec",
 			code: `package main

decorator/fragger-generated.go

@@ -968,6 +968,8 @@ func (f *Fragger) ProcessNode(n ast.Node) {
 
 	case *ast.Package:
 
+		// Map: Imports
+
 		// Map: Files
 		for _, v := range n.Files {
 			f.ProcessNode(v)