termui-tree/tree.go

package main

import (
	"image"

	"github.com/gizak/termui/v3"
)

const indentWidth = 4

// TreeWidget is a renderable tree
type TreeWidget struct {
	termui.Block
	rootNodes    []*Node
	selectedNode *Node
}

// Node represents an item in the tree widget
type Node struct {
	Text     string
	children []*Node
	tree     *TreeWidget
	parent   *Node
	next     *Node
	prev     *Node
}

// NewTreeWidget makes a new tree widget
func NewTreeWidget(rootNodes []*Node) *TreeWidget {
	tree := &TreeWidget{
		Block: *termui.NewBlock(),
	}
	tree.selectedNode = rootNodes[0]
	for _, rootNode := range rootNodes {
		tree.AddRootNode(rootNode)
	}

	return tree
}

// Draw draws the tree widget
func (tree *TreeWidget) Draw(buffer *termui.Buffer) {
	tree.Block.Draw(buffer)
	row := tree.Block.Inner.Min.Y
	for _, item := range tree.rootNodes {
		row = tree.renderSingleTree(item, buffer, row)
	}
}

// AddRootNode adds a new top-level node to the tree
func (tree *TreeWidget) AddRootNode(node *Node) {
	if len(tree.rootNodes) > 0 {
		node.prev = tree.rootNodes[len(tree.rootNodes)-1]
		node.prev.next = node
	} else if len(tree.rootNodes) == 0 {
		// If this is our first node, then we need to mark this as the current node
		tree.selectedNode = node
	}
	tree.rootNodes = append(tree.rootNodes, node)
	node.setTree(tree)
}

// SelectPrev marks the previous node as the currently selected node
func (tree *TreeWidget) SelectPrev() {
	if tree.selectedNode.prev != nil && len(tree.selectedNode.prev.children) == 0 {
		// If there's a previous node with no children, we can go straight to it.
		tree.selectedNode = tree.selectedNode.prev
	} else if tree.selectedNode.prev == nil && tree.selectedNode.parent != nil {
		// If we can go right up to the next parent node, we're good - we go there.
		tree.selectedNode = tree.selectedNode.parent
	} else if tree.selectedNode.prev != nil {
		// If we have a node before us, we're going to want to go to the absolute bottom of it.
		cursorNode := tree.selectedNode.prev
		for len(cursorNode.children) > 0 {
			cursorNode = cursorNode.children[len(cursorNode.children)-1]
		}
		tree.selectedNode = cursorNode
	}
}

// SelectNext marks the next node as the currently selected node
func (tree *TreeWidget) SelectNext() {
	if len(tree.selectedNode.children) > 0 {
		tree.selectedNode = tree.selectedNode.children[0]
	} else if tree.selectedNode.next != nil {
		tree.selectedNode = tree.selectedNode.next
	} else if tree.selectedNode.next == nil && tree.selectedNode.parent != nil {
		cursorNode := tree.selectedNode.parent
		// Go up to the next adjacent node to the first parent node with another node adjacent to it.
		for cursorNode != nil {
			if cursorNode.next != nil {
				tree.selectedNode = cursorNode.next
				return
			}
			cursorNode = cursorNode.parent
		}
	}
	// If all else fails, we must be at the end of the list, and there's nothing else to select.
}

// renderSingleTree renders a node and all its descendents
// Returns the number of rows needed to render it (inculding wrapping)
func (tree *TreeWidget) renderSingleTree(root *Node, buffer *termui.Buffer, startRow int) int {
	baseWidth := tree.Inner.Max.X - tree.Inner.Min.X

	depths := map[*Node]int{root: 0}
	toVisit := []*Node{root}
	row := startRow
	for len(toVisit) > 0 {
		var node *Node
		toVisit, node = toVisit[:len(toVisit)-1], toVisit[len(toVisit)-1]
		currentDepth := depths[node]
		// Add the children to explore next in reverse order.
		// The next off the stack should be the direct descendent of the parent.
		for i := len(node.children) - 1; i >= 0; i-- {
			child := node.children[i]
			toVisit = append(toVisit, child)
			depths[child] = currentDepth + 1
		}

		// Render this level of tree
		width := baseWidth - currentDepth*indentWidth
		startingCol := tree.Inner.Min.X + currentDepth*indentWidth
		row += node.render(buffer, row, startingCol, width)
	}

	return row
}

// render renders this node in the term ui, and return the number of rows that were required to render it
func (node *Node) render(buffer *termui.Buffer, row, col, width int) int {
	style := termui.Style{
		Fg: termui.ColorRed,
		Bg: termui.ColorClear,
	}
	if node.tree.selectedNode == node {
		style.Fg, style.Bg = style.Bg, style.Fg
	}

	listItem := termui.ParseStyles(node.Text, style)
	wrappedItem := termui.WrapCells(listItem, uint(width))
	x := col
	rowOffset := 0
	for _, cell := range wrappedItem {
		// If we're wrapped, we don't need to render anything, just go to the next line
		if cell.Rune == '\n' {
			rowOffset++
			x = col
		} else {
			buffer.SetCell(cell, image.Point{x, row + rowOffset})
			x++
		}
	}

	return rowOffset + 1
}

// AddChild marks the given node as a child, and itself as the child node's parent
func (node *Node) AddChild(child *Node) {
	child.parent = node
	if node.children == nil {
		node.children = []*Node{child}
	} else {
		node.children = append(node.children, child)
	}
	child.tree = node.tree
	if len(node.children) > 1 {
		lastNode := node.children[len(node.children)-2]
		lastNode.next = child
		child.prev = lastNode
	}
}

func (node *Node) setTree(tree *TreeWidget) {
	node.tree = tree
	for _, child := range node.children {
		child.setTree(tree)
	}
}