exercism/go/pov/tree/pov.go

package pov

const testVersion = 2

type Graph struct {
	leaves []Node
	Iters  int
}

func New() *Graph {
	return new(Graph)
}

func (g *Graph) AddNode(lbl string) {
	// Make sure that g doesn't already have node lbl
	if g.getNode(lbl) != nil {
		return
	}
	g.leaves = append(g.leaves, Node{label: lbl})
}

func (g *Graph) AddArc(fr, to string) {
	var frN *Node
	if frN = g.getNode(fr); frN == nil {
		// frN doesn't exist yet... Create it and
		// drop
		frN = &Node{label: fr}
	}
	frN.AddNode(to)
}

func (g *Graph) ArcList() []string {
	var arcList []string
	for i := range g.leaves {
		arcList = append(arcList, g.leaves[i].ArcList()...)
	}
	return arcList
}

func (g *Graph) GetPath(fr, to string) []string {
	if frN := g.getNode(fr); frN != nil {
		return frN.GetPath(to)
	}
	return []string{}
}

func (g *Graph) getNode(lbl string) *Node {
	for i := range g.leaves {
		if r := g.leaves[i].getNode(lbl); r != nil {
			return r
		}
	}
	return nil
}

// Return a graph from newRoot's POV
func (g *Graph) ChangeRoot(oldRoot, newRoot string) *Graph {
	ret := g
	pth := ret.GetPath(oldRoot, newRoot)
	for oldRoot != newRoot {
		// Swap the first two in pth
		if len(pth) >= 2 {
			n := ret.getNode(pth[1])
			if n == nil {
				return ret
			}
			n = n.copy()
			p := ret.getNode(pth[0])
			if p == nil {
				return ret
			}
			p = p.copy()
			if n != nil && p != nil {
				p.removeLeaf(n.label)
				n.addWholeNode(p)
				if p.parent != nil {
					n.parent = p.parent
					p.parent = nil
				} else {
					n.parent = nil
				}
			}
			ret = New()
			ret.addWholeNode(n)
			oldRoot = n.label
		}
		pth = ret.GetPath(oldRoot, newRoot)
	}
	return ret
}

/*
// Return a graph from newRoot's POV
func (g *Graph) ChangeRoot(oldRoot, newRoot string) *Graph {
	if oldRoot == newRoot {
		return g
	}
	ret := New()
	pth := g.GetPath(oldRoot, newRoot)
	fmt.Println(pth)
	// Swap the first two in pth
	if len(pth) >= 2 {
		fmt.Println("Swap ", pth[0], "<>", pth[1])
		n := g.getNode(pth[1])
		p := g.getNode(pth[0])
		if n != nil && p != nil {
			p.removeLeaf(n.label)
			n.addWholeNode(p)
			if p.parent != nil {
				n.parent = p.parent
				p.parent = nil
			} else {
				n.parent = nil
			}
		}
		ret.addWholeNode(n)
	}
	PrintGraph(ret)
	if pth[1] != newRoot {
		fmt.Println("Recurse:", pth[1], "<>", newRoot)
		return ret
		//return ret.ChangeRoot(pth[1], newRoot)
	} else {
		return ret
	}
}
*/

func (g *Graph) addWholeNode(n *Node) {
	g.leaves = append(g.leaves, *n)
}

type Node struct {
	label  string
	leaves []Node
	parent *Node
}

func (n *Node) copy() *Node {
	// Returns a new node that is a copy of this node
	ret := Node{label: n.label, parent: n.parent}
	for _, k := range n.leaves {
		cp := k.copy()
		ret.leaves = append(ret.leaves, *cp)
	}
	return &ret
}

func (n *Node) AddNode(lbl string) {
	nn := Node{label: lbl, parent: n}
	n.leaves = append(n.leaves, nn)
}

func (n *Node) addWholeNode(nd *Node) {
	n.leaves = append(n.leaves, *nd)
}

func (n *Node) removeLeaf(lbl string) {
	for i, k := range n.leaves {
		if k.label == lbl {
			copy(n.leaves[i:], n.leaves[i+1:])
			n.leaves[len(n.leaves)-1] = Node{}
			n.leaves = n.leaves[:len(n.leaves)-1]
		}
	}
}

func (n *Node) ArcList() []string {
	var ret []string
	for i := range n.leaves {
		// Append this arc
		ret = append(ret, n.label+" -> "+n.leaves[i].label)
		// Append all sub-arcs
		ret = append(ret, n.leaves[i].ArcList()...)
	}
	return ret
}

// GetPath returns a string slice of all leaves leading to 'to'
func (n *Node) GetPath(to string) []string {
	ret := []string{n.label}
	if n.label == to {
		return ret
	}
	for i := range n.leaves {
		if toN := n.leaves[i].getNode(to); toN != nil {
			// This is the leaf to follow
			return append(ret, n.leaves[i].GetPath(to)...)
		}
	}
	// Couldn't get there
	return []string{}
}

// Returns the node with label lbl or nil
func (n *Node) getNode(lbl string) *Node {
	if n.label == lbl {
		return n
	}
	for i := range n.leaves {
		if r := n.leaves[i].getNode(lbl); r != nil {
			return r
		}
	}
	return nil
}