commonMain.Node.kt Maven / Gradle / Ivy

Go to download
package com.juul.krayon.hierarchy

public class Node internal constructor(
    /** Data associated with this node. */
    public val data: T,

    layout: L,

    /** The parent node, if any. */
    public val parent: Node? = null,
) {

    /** Layout value. Null when initially created, but can be changed by using a tree-map or similar.  */
    public var layout: L = layout
        internal set

    /** Child nodes. If this is a leaf, then this list is empty. */
    public var children: List> = emptyList()
        internal set

    /** Weight associated with this item. Note that this defaults to 0, and must be explicitly set via [sum] or [count]. */
    @PublishedApi
    internal var weight: Float = 0f
}

/** Number of parents before hitting the root [Node]. If this is the root, then this is 0. */
public val  Node.depth: Int
    get() = if (parent == null) 0 else parent.depth + 1

/** Maximum number of children before hitting a leaf [Node]. If this is a leaf, then this is 0. */
public val  Node.height: Int
    get() = children.maxOfOrNull { it.height + 1 } ?: 0

/** Returns `true` if this node has no children. */
public val  Node.isLeaf: Boolean
    get() = children.isEmpty()

/** Returns ancestor nodes, starting with `this` and then following the [Node.parent] chain. */
public fun  Node.ancestors(): Sequence> = sequence {
    var current: Node? = this@ancestors
    while (current != null) {
        yield(current)
        current = parent
    }
}

/** Returns descendant nodes, starting with `this` and then following the [Node.children] chain in a breadth-first traversal. */
public fun  Node.traverseBreadthFirst(): Sequence> =
    sequence {
        val remaining = ArrayDeque(listOf(this@traverseBreadthFirst))
        while (remaining.isNotEmpty()) {
            val current = remaining.removeFirst()
            yield(current)
            current.children.forEach(remaining::addLast)
        }
    }

/** Returns descendant nodes, starting with `this` and then following the [Node.children] chain in a depth-first traversal. */
public fun  Node.traversePreOrder(): Sequence> =
    sequence {
        yield(this@traversePreOrder)
        for (child in children) {
            yieldAll(child.traversePreOrder())
        }
    }

/** Returns descendant nodes, starting with the [Node.children] chain in a depth-first traversal, and ending with `this`. */
public fun  Node.traversePostOrder(): Sequence> =
    sequence {
        for (child in children) {
            yieldAll(child.traversePostOrder())
        }
        yield(this@traversePostOrder)
    }

public inline fun  Node.sum(
    crossinline value: (T) -> Float,
): Node = eachAfter { node -> node.weight = value(node.data) + node.children.sumOf { it.weight } }

public fun  Node.count(): Node = sum { 1f }

public fun  Node.sort(
    comparator: Comparator,
): Node = eachBefore { node ->
    node.children = node.children.sortedWith { left, right ->
        comparator.compare(left.data, right.data)
    }
}

public inline fun  Node.each(
    crossinline action: (Node) -> Unit,
): Node = apply { traverseBreadthFirst().forEach(action) }

public inline fun  Node.eachAfter(
    crossinline action: (Node) -> Unit,
): Node = apply { traversePostOrder().forEach(action) }

public inline fun  Node.eachBefore(
    crossinline action: (Node) -> Unit,
): Node = apply { traversePreOrder().forEach(action) }

public inline fun  Node.eachIndexed(
    crossinline action: (Int, Node) -> Unit,
): Node = apply { traverseBreadthFirst().forEachIndexed(action) }

public inline fun  Node.eachAfterIndexed(
    crossinline action: (Int, Node) -> Unit,
): Node = apply { traversePostOrder().forEachIndexed(action) }

public inline fun  Node.eachBeforeIndexed(
    crossinline action: (Int, Node) -> Unit,
): Node = apply { traversePreOrder().forEachIndexed(action) }

/**
 * Returns a breadth-first traversal of nodes, removing those with null data, and pairing the data to the layout.
 *
 * Especially useful with [flatHierarchy] after performing a layout, before feeding into selection data.
 */
public fun  Node.removeHierarchy(): Sequence> =
    traverseBreadthFirst()
        .filter { it.data != null }
        .map { checkNotNull(it.data) to it.layout }