Re: [PATCH v8 1/2] kho: Adopt radix tree for preserved memory tracking

[Mike Rapoport <rppt@kernel.org>]

On Thu, Jan 29, 2026 at 04:57:38PM -0800, Jason Miu wrote:
> Introduce a radix tree implementation for tracking preserved memory
> pages and switch the KHO memory tracking mechanism to use it. This
> lays the groundwork for a stateless KHO implementation that eliminates
> the need for serialization and the associated "finalize" state.
> 
> This patch introduces the core radix tree data structures and
> constants to the KHO ABI. It adds the radix tree node and leaf
> structures, along with documentation for the radix tree key encoding
> scheme that combines a page's physical address and order.
> 
> To support broader use by other kernel subsystems, such as hugetlb
> preservation, the core radix tree manipulation functions are exported
> as a public API.
> 
> The xarray-based memory tracking is replaced with this new radix tree
> implementation. The core KHO preservation and unpreservation functions
> are wired up to use the radix tree helpers. On boot, the second kernel
> restores the preserved memory map by walking the radix tree whose root
> physical address is passed via the FDT.
> 
> The ABI `compatible` version is bumped to "kho-v2" to reflect the
> structural changes in the preserved memory map and sub-FDT property
> names. This includes renaming "fdt" to "preserved-data" to better
> reflect that preserved state may use formats other than FDT.
> 
> Signed-off-by: Jason Miu <jasonmiu@google.com>

With minor comments below

Reviewed-by: Mike Rapoport (Microsoft) <rppt@kernel.org>

> ---
...

> +static int __kho_radix_walk_tree(struct kho_radix_node *root,
> +				 unsigned int level, unsigned long start,
> +				 kho_radix_tree_walk_callback_t cb)
> +{
> +	struct kho_radix_node *node;
> +	struct kho_radix_leaf *leaf;
> +	unsigned long key, i;
> +	unsigned int shift;
> +	int err;
> +
> +	for (i = 0; i < PAGE_SIZE / sizeof(phys_addr_t); i++) {
> +		if (!root->table[i])
> +			continue;
> +
> +		shift = ((level - 1) * KHO_TABLE_SIZE_LOG2) +
> +			KHO_BITMAP_SIZE_LOG2;
> +		key = start | (i << shift);
> +
> +		node = phys_to_virt(root->table[i]);
> +
> +		if (level == 1) {
> +			/*
> +			 * we are at level 1,
> +			 * node is pointing to the level 0 bitmap.
> +			 */
> +			leaf = (struct kho_radix_leaf *)node;
> +			err = kho_radix_walk_leaf(leaf, key, cb);
> +			if (err)
> +				return err;
> +		} else {
> +			err  = __kho_radix_walk_tree(node, level - 1,
> +						     key, cb);
> +			if (err)
> +				return err;

Nit: if (err) can be moved outside if (level == 1)

> +		}
> +	}
> +
> +	return 0;
> +}
> +
> +/**
> + * kho_radix_walk_tree - Traverses the radix tree and calls a callback for each preserved page.
> + * @tree: A pointer to the KHO radix tree to walk.
> + * @cb: A callback function of type kho_radix_tree_walk_callback_t that will be
> + *      invoked for each preserved page found in the tree. The callback receives
> + *      the physical address and order of the preserved page.
> + *
> + * This function walks the radix tree, searching from the specified top level
> + * (@level) down to the lowest level (level 0). For each preserved page found,

Hmm, why do we need @level here? 
Or it rather remained from the older versions?

> + * it invokes the provided callback, passing the page's physical address and
> + * order.
> + *
> + * Return: 0 if the walk completed the specified tree, or the non-zero return
> + *         value from the callback that stopped the walk.
> + */
> +int kho_radix_walk_tree(struct kho_radix_tree *tree,
> +			kho_radix_tree_walk_callback_t cb)
> +{
> +	if (WARN_ON_ONCE(!tree->root))
> +		return -EINVAL;
> +
> +	guard(mutex)(&tree->lock);
> +
> +	return __kho_radix_walk_tree(tree->root, KHO_TREE_MAX_DEPTH - 1, 0, cb);
> +}
> +EXPORT_SYMBOL_GPL(kho_radix_walk_tree);

-- 
Sincerely yours,
Mike.