[RFC PATCH v4 2/7] libbpf: Add BTF permutation support for type reordering

[Donglin Peng <dolinux.peng@gmail.com>]

From: pengdonglin <pengdonglin@xiaomi.com>

Introduce btf__permute() API to allow in-place rearrangement of BTF types.
This function reorganizes BTF type order according to a provided array of
type IDs, updating all type references to maintain consistency.

The permutation process involves:
1. Shuffling types into new order based on the provided ID mapping
2. Remapping all type ID references to point to new locations
3. Handling BTF extension data if provided via options

This is particularly useful for optimizing type locality after BTF
deduplication or for meeting specific layout requirements in specialized
use cases.

Cc: Eduard Zingerman <eddyz87@gmail.com>
Cc: Alexei Starovoitov <ast@kernel.org>
Cc: Andrii Nakryiko <andrii.nakryiko@gmail.com>
Cc: Alan Maguire <alan.maguire@oracle.com>
Cc: Song Liu <song@kernel.org>
Signed-off-by: pengdonglin <pengdonglin@xiaomi.com>
Signed-off-by: Donglin Peng <dolinux.peng@gmail.com>
---
 tools/lib/bpf/btf.c      | 161 +++++++++++++++++++++++++++++++++++++++
 tools/lib/bpf/btf.h      |  34 +++++++++
 tools/lib/bpf/libbpf.map |   1 +
 3 files changed, 196 insertions(+)

diff --git a/tools/lib/bpf/btf.c b/tools/lib/bpf/btf.c
index 5e1c09b5dce8..3bc03f7fe31f 100644
--- a/tools/lib/bpf/btf.c
+++ b/tools/lib/bpf/btf.c
@@ -5830,3 +5830,164 @@ int btf__relocate(struct btf *btf, const struct btf *base_btf)
 		btf->owns_base = false;
 	return libbpf_err(err);
 }
+
+struct btf_permute {
+	/* .BTF section to be permuted in-place */
+	struct btf *btf;
+	struct btf_ext *btf_ext;
+	/* Array of type IDs used for permutation. The array length must equal
+	 * the number of types in the BTF being permuted, excluding the special
+	 * void type at ID 0. For split BTF, the length corresponds to the
+	 * number of types added on top of the base BTF.
+	 */
+	__u32 *ids;
+	/* Array of type IDs used to map from original type ID to a new permuted
+	 * type ID, its length equals to the above ids */
+	__u32 *map;
+};
+
+static int btf_permute_shuffle_types(struct btf_permute *p);
+static int btf_permute_remap_types(struct btf_permute *p);
+static int btf_permute_remap_type_id(__u32 *type_id, void *ctx);
+
+int btf__permute(struct btf *btf, __u32 *ids, const struct btf_permute_opts *opts)
+{
+	struct btf_permute p;
+	int i, err = 0;
+	__u32 *map = NULL;
+
+	if (!OPTS_VALID(opts, btf_permute_opts) || !ids)
+		return libbpf_err(-EINVAL);
+
+	map = calloc(btf->nr_types, sizeof(*map));
+	if (!map) {
+		err = -ENOMEM;
+		goto done;
+	}
+
+	for (i = 0; i < btf->nr_types; i++)
+		map[i] = BTF_UNPROCESSED_ID;
+
+	p.btf = btf;
+	p.btf_ext = OPTS_GET(opts, btf_ext, NULL);
+	p.ids = ids;
+	p.map = map;
+
+	if (btf_ensure_modifiable(btf)) {
+		err = -ENOMEM;
+		goto done;
+	}
+	err = btf_permute_shuffle_types(&p);
+	if (err < 0) {
+		pr_debug("btf_permute_shuffle_types failed: %s\n", errstr(err));
+		goto done;
+	}
+	err = btf_permute_remap_types(&p);
+	if (err < 0) {
+		pr_debug("btf_permute_remap_types failed: %s\n", errstr(err));
+		goto done;
+	}
+
+done:
+	free(map);
+	return libbpf_err(err);
+}
+
+/* Shuffle BTF types.
+ *
+ * Rearranges types according to the permutation map in p->ids. The p->map
+ * array stores the mapping from original type IDs to new shuffled IDs,
+ * which is used in the next phase to update type references.
+ *
+ * Validates that all IDs in the permutation array are valid and unique.
+ */
+static int btf_permute_shuffle_types(struct btf_permute *p)
+{
+	struct btf *btf = p->btf;
+	const struct btf_type *t;
+	__u32 *new_offs = NULL, *map;
+	void *nt, *new_types = NULL;
+	int i, id, len, err;
+
+	new_offs = calloc(btf->nr_types, sizeof(*new_offs));
+	new_types = calloc(btf->hdr->type_len, 1);
+	if (!new_offs || !new_types) {
+		err = -ENOMEM;
+		goto out_err;
+	}
+
+	nt = new_types;
+	for (i = 0; i < btf->nr_types; i++) {
+		id = p->ids[i];
+		/* type IDs from base_btf and the VOID type are not allowed */
+		if (id < btf->start_id) {
+			err = -EINVAL;
+			goto out_err;
+		}
+		/* must be a valid type ID */
+		t = btf__type_by_id(btf, id);
+		if (!t) {
+			err = -EINVAL;
+			goto out_err;
+		}
+		map = &p->map[id - btf->start_id];
+		/* duplicate type IDs are not allowed */
+		if (*map != BTF_UNPROCESSED_ID) {
+			err = -EINVAL;
+			goto out_err;
+		}
+		len = btf_type_size(t);
+		memcpy(nt, t, len);
+		new_offs[i] = nt - new_types;
+		*map = btf->start_id + i;
+		nt += len;
+	}
+
+	free(btf->types_data);
+	free(btf->type_offs);
+	btf->types_data = new_types;
+	btf->type_offs = new_offs;
+	return 0;
+
+out_err:
+	free(new_offs);
+	free(new_types);
+	return err;
+}
+
+/* Callback function to remap individual type ID references
+ *
+ * This callback is invoked by btf_remap_types() for each type ID reference
+ * found in the BTF data. It updates the reference to point to the new
+ * permuted type ID using the mapping table.
+ */
+static int btf_permute_remap_type_id(__u32 *type_id, void *ctx)
+{
+	struct btf_permute *p = ctx;
+	__u32 new_type_id = *type_id;
+
+	/* skip references that point into the base BTF */
+	if (new_type_id < p->btf->start_id)
+		return 0;
+
+	new_type_id = p->map[*type_id - p->btf->start_id];
+	if (new_type_id > BTF_MAX_NR_TYPES)
+		return -EINVAL;
+
+	*type_id = new_type_id;
+	return 0;
+}
+
+/* Remap referenced type IDs into permuted type IDs.
+ *
+ * After BTF types are permuted, their final type IDs may differ from original
+ * ones. The map from original to a corresponding permuted type ID is stored
+ * in btf_permute->map and is populated during shuffle phase. During remapping
+ * phase we are rewriting all type IDs  referenced from any BTF type (e.g.,
+ * struct fields, func proto args, etc) to their final deduped type IDs.
+ */
+static int btf_permute_remap_types(struct btf_permute *p)
+{
+	return btf_remap_types(p->btf, p->btf_ext, btf_permute_remap_type_id, p);
+}
+
diff --git a/tools/lib/bpf/btf.h b/tools/lib/bpf/btf.h
index ccfd905f03df..441f6445d762 100644
--- a/tools/lib/bpf/btf.h
+++ b/tools/lib/bpf/btf.h
@@ -273,6 +273,40 @@ LIBBPF_API int btf__dedup(struct btf *btf, const struct btf_dedup_opts *opts);
  */
 LIBBPF_API int btf__relocate(struct btf *btf, const struct btf *base_btf);
 
+struct btf_permute_opts {
+	size_t sz;
+	/* optional .BTF.ext info along the main BTF info */
+	struct btf_ext *btf_ext;
+	size_t :0;
+};
+#define btf_permute_opts__last_field btf_ext
+
+/**
+ * @brief **btf__permute()** rearranges BTF types in-place according to specified mapping
+ * @param btf BTF object to permute
+ * @param ids Array defining new type order. Must contain exactly btf->nr_types elements,
+ *        each being a valid type ID in range [btf->start_id, btf->start_id + btf->nr_types - 1]
+ * @param opts Optional parameters, including BTF extension data for reference updates
+ * @return 0 on success, negative error code on failure
+ *
+ * **btf__permute()** performs an in-place permutation of BTF types, rearranging them
+ * according to the order specified in @p ids array. After reordering, all type references
+ * within the BTF data and optional BTF extension are updated to maintain consistency.
+ *
+ * The permutation process consists of two phases:
+ * 1. Type shuffling: Physical reordering of type data in memory
+ * 2. Reference remapping: Updating all type ID references to new locations
+ *
+ * This is particularly useful for optimizing type locality after BTF deduplication
+ * or for meeting specific layout requirements in specialized use cases.
+ *
+ * On error, negative error code is returned and errno is set appropriately.
+ * Common error codes include:
+ *   - -EINVAL: Invalid parameters or invalid ID mapping (e.g., duplicate IDs, out-of-range IDs)
+ *   - -ENOMEM: Memory allocation failure during permutation process
+ */
+LIBBPF_API int btf__permute(struct btf *btf, __u32 *ids, const struct btf_permute_opts *opts);
+
 struct btf_dump;
 
 struct btf_dump_opts {
diff --git a/tools/lib/bpf/libbpf.map b/tools/lib/bpf/libbpf.map
index 8ed8749907d4..b778e5a5d0a8 100644
--- a/tools/lib/bpf/libbpf.map
+++ b/tools/lib/bpf/libbpf.map
@@ -451,4 +451,5 @@ LIBBPF_1.7.0 {
 	global:
 		bpf_map__set_exclusive_program;
 		bpf_map__exclusive_program;
+		btf__permute;
 } LIBBPF_1.6.0;
-- 
2.34.1