如何在golang ast遍历中从子节点检索父节点？

Question

Here is test file

func demo()  {
    name:=xxx()
    fmt.Println(name)
}

And my ast traversal code

ast.Inspect(f, func(node ast.Node) bool {
    assign, ok := node.(*ast.AssignStmt) // find -> name:=xxx()
    if !ok {
        return true
    }
    if assign == nil {
        return true
    }
    var expr = assign.Lhs[0]
    fmt.Println(expr) // find -> name
    fmt.Println(nodeString(assign,pass.Fset))
    return true
})

I find I have to travel ast from up to down like this. Find AssignStmt and then get AssignStmt->Lhs

But I need to find ast.Ident(name) first and then to find whether its parent is AssignStmt in some cases.

I am not sure is it possible that I can go from bottom to top.

Answer 1

Manage a stack of ancestor nodes as the inspector is called:

var stack []ast.Node
ast.Inspect(f, func(n ast.Node) bool {

    // Look for the identifier.

    if n, ok := n.(*ast.Ident); ok {
        if n.Name == "name" {

            // Parent is stack[len(stack)-1]

            fmt.Println("found name")
            for _, n := range stack {
                fmt.Printf("%T\n", n)
            }
        }
    }

    // Manage the stack. Inspect calls a function like this:
    //   f(node)
    //   for each child {
    //      f(child) // and recursively for child's children
    //   }
    //   f(nil)
    if n == nil {
        // Done with node's children. Pop.
        stack = stack[:len(stack)-1]
    } else {
        // Push the current node for children.
        stack = append(stack, n)
    }

    return true
})

Run it on the playground .

Answer 2

I came across similar problem that I needed to find the index of the child amongst its siblings as well as the parent node of it. I could not find closer problem so I am posting here.

I isolated the logic in Cerise's answer into a function and moved the answer specific parts to outside of function. So more people can leverage this generic form and save time.

Logic for finding the index is similar to finding parent. Another array for tracing which -index sequence- is followed to reach current node is enough. At each call where n is nil we remove the last item from sequence and increment the new last item.

func FindParentNodeAndChildIndex(root ast.Node, child ast.Node) (parent ast.Node, childIndex int) {
    var (
        parentTrace     []ast.Node
        childIndexTrace []int
        found           = false
    )
    var (
        updateParentTrace = func(n ast.Node) {
            if n != nil {
                parentTrace = append(parentTrace, n)
            } else {
                parentTrace = parentTrace[:len(parentTrace)-1]
            }
        }
        updateChildIndexTrace = func(n ast.Node) {
            if n != nil {
                childIndexTrace = append(childIndexTrace, 0)
            } else {
                l := len(childIndexTrace)
                childIndexTrace = childIndexTrace[:l-1]
                childIndexTrace[l-2]++
            }
        }
    )
    ast.Inspect(root, func(n ast.Node) bool {
        if !found {
            updateParentTrace(n)
            updateChildIndexTrace(n)
        }
        if n != nil && n == child {
            found = true
        }
        return !found
    })
    if found {
        return parentTrace[len(parentTrace)-2], childIndexTrace[len(childIndexTrace)-2]
    }
    return nil, -1
}

It looks like there are more alternative methods for finding parent then the date question answered. One is astutil.Apply() another is inspector.WithStack()

astutil.Apply() is used with two callbacks one fired before and another after visiting each node's children; and accepts a Cursor instance that holds details about visited node's Index and Parent.

For WithStack there is a detailed example