From 0c07ff6ac8434afc391695bd3687355f7b4436f2 Mon Sep 17 00:00:00 2001
From: Austin Schuh <austin.linux@gmail.com>
Date: Mon, 20 Jan 2025 22:06:11 -0800
Subject: [PATCH] Support zstd compressed tegraflash images

zstd is a lot faster when compressing and decompressing, which helps
with iteration speed.

Use the CONVERSIONTYPES machinery builtin rather than re-implementing it
ourselves here.

Signed-off-by: Austin Schuh <austin.linux@gmail.com>
Signed-off-by: Matt Madison <matt@madison.systems>
---
 classes-recipe/image_types_tegra.bbclass | 29 +++++-------------------
 conf/machine/include/tegra-common.inc    |  2 +-
 2 files changed, 7 insertions(+), 24 deletions(-)

diff --git a/classes-recipe/image_types_tegra.bbclass b/classes-recipe/image_types_tegra.bbclass
index 5057f596..2ff6830f 100644
--- a/classes-recipe/image_types_tegra.bbclass
+++ b/classes-recipe/image_types_tegra.bbclass
@@ -4,7 +4,7 @@ TEGRA_UEFI_SIGNING_CLASS ??= "tegra-uefi-signing"
 inherit ${TEGRA_UEFI_SIGNING_CLASS}
 TEGRA_UEFI_USE_SIGNED_FILES ??= "false"
 
-IMAGE_TYPES += "tegraflash"
+IMAGE_TYPES += "tegraflash.tar"
 
 IMAGE_ROOTFS_ALIGNMENT ?= "4"
 
@@ -103,17 +103,8 @@ TOSIMGFILENAME:tegra234 = "tos-optee_t234.img"
 BUP_PAYLOAD_DIR = "payloads_t${@d.getVar('NVIDIA_CHIP')[2:]}x"
 FLASHTOOLS_DIR = "tegra-flash"
 
-TEGRAFLASH_PACKAGE_FORMAT ??= "tar"
-TEGRAFLASH_PACKAGE_FORMATS = "tar zip"
 TEGRAFLASH_SDCARD_SIZE ??= "16G"
 
-python() {
-    pt = d.getVar('TEGRAFLASH_PACKAGE_FORMAT')
-    valid_types = d.getVar('TEGRAFLASH_PACKAGE_FORMATS').split()
-    if pt not in valid_types:
-        bb.fatal("TEGRAFLASH_PACKAGE_FORMAT %s not recognized, supported types are: %s" % (pt, ', '.join(valid_types)))
-}
-
 # Override this function if you need to add
 # customization after the default files are
 # copied/symlinked into the working directory
@@ -174,15 +165,8 @@ tegraflash_custom_post() {
 }
 
 tegraflash_finalize_pkg() {
-    if [ "${TEGRAFLASH_PACKAGE_FORMAT}" = "zip" ]; then
-        rm -f ${IMGDEPLOYDIR}/${IMAGE_NAME}.tegraflash.zip
-        zip -r ${IMGDEPLOYDIR}/${IMAGE_NAME}.tegraflash.zip .
-        ln -sf ${IMAGE_NAME}.tegraflash.zip ${IMGDEPLOYDIR}/${IMAGE_LINK_NAME}.tegraflash.zip
-    else
-        rm -f ${IMGDEPLOYDIR}/${IMAGE_NAME}.tegraflash.tar.gz
-        ${IMAGE_CMD_TAR} --sparse --numeric-owner --transform="s,^\./,," -cf- . | gzip -f -9 -n -c --rsyncable > ${IMGDEPLOYDIR}/${IMAGE_NAME}.tegraflash.tar.gz
-        ln -sf ${IMAGE_NAME}.tegraflash.tar.gz ${IMGDEPLOYDIR}/${IMAGE_LINK_NAME}.tegraflash.tar.gz
-    fi
+    rm -f ${IMGDEPLOYDIR}/${IMAGE_NAME}.tegraflash.tar
+    ${IMAGE_CMD_TAR} --sparse --numeric-owner --transform="s,^\./,," -cf ${IMGDEPLOYDIR}/${IMAGE_NAME}.tegraflash.tar .
 }
 
 tegraflash_create_flash_config() {
@@ -449,9 +433,8 @@ EOF
     chmod +x $outfile
 }
 
-IMAGE_CMD:tegraflash = "create_tegraflash_pkg"
-TEGRAFLASH_PKG_DEPENDS = "${@'zip-native:do_populate_sysroot' if d.getVar('TEGRAFLASH_PACKAGE_FORMAT') == 'zip' else '${CONVERSION_DEPENDS_gz}:do_populate_sysroot'}"
-do_image_tegraflash[depends] += "${TEGRAFLASH_PKG_DEPENDS} dtc-native:do_populate_sysroot coreutils-native:do_populate_sysroot \
+IMAGE_CMD:tegraflash.tar = "create_tegraflash_pkg"
+do_image_tegraflash_tar[depends] += "dtc-native:do_populate_sysroot coreutils-native:do_populate_sysroot \
                                  tegra-flashtools-native:do_populate_sysroot gptfdisk-native:do_populate_sysroot \
                                  tegra-bootfiles:do_populate_sysroot tegra-bootfiles:do_populate_lic \
                                  virtual/kernel:do_deploy \
@@ -459,4 +442,4 @@ do_image_tegraflash[depends] += "${TEGRAFLASH_PKG_DEPENDS} dtc-native:do_populat
                                  ${@'${TEGRA_ESP_IMAGE}:do_image_complete' if d.getVar('TEGRA_ESP_IMAGE') != '' else  ''} \
                                  virtual/bootloader:do_deploy virtual/secure-os:do_deploy ${TEGRA_SIGNING_EXTRA_DEPS} ${DTB_EXTRA_DEPS} \
                                  ${@'${TEGRAFLASH_INITRD_FLASH_IMAGE}:do_image_complete' if d.getVar('TEGRAFLASH_INITRD_FLASH_IMAGE') != '' else ''}"
-IMAGE_TYPEDEP:tegraflash += "${IMAGE_TEGRAFLASH_FS_TYPE}"
+IMAGE_TYPEDEP:tegraflash.tar += "${IMAGE_TEGRAFLASH_FS_TYPE}"
diff --git a/conf/machine/include/tegra-common.inc b/conf/machine/include/tegra-common.inc
index 9707a909..7318b8c4 100644
--- a/conf/machine/include/tegra-common.inc
+++ b/conf/machine/include/tegra-common.inc
@@ -67,7 +67,7 @@ IMAGE_ROOTFS_ALIGNMENT ?= "4"
 TEGRA_BLBLOCKSIZE ?= "${@int(d.getVar('IMAGE_ROOTFS_ALIGNMENT')) * 1024}"
 EXTRA_IMAGECMD:ext4 ?= "-i 4096 -b 4096"
 IMAGE_CLASSES += "image_types_tegra"
-IMAGE_FSTYPES += "tegraflash"
+IMAGE_FSTYPES += "tegraflash.tar.zst"
 
 INITRAMFS_IMAGE ?= "tegra-minimal-initramfs"
 INITRAMFS_IMAGE_BUNDLE ?= "0"