TheRock/rocm-systems/shared/kpack/python/rocm_kpack/elf/zero_page.py

"""
Zero-page optimization for ELF binaries.

This module handles the removal of content from sections while maintaining
ELF validity. The key operation is "zero-paging" - converting file-backed
content to memory-only (NOBITS) so the file can be smaller.

The conservative zero-page algorithm:
1. Find page-aligned region within the section
2. Remove that content from the file
3. Split the containing PT_LOAD into multiple segments
4. Adjust all offsets for the removed bytes
5. Mark the section as NOBITS

Key invariants:
- Only full pages can be removed (partial pages must be preserved)
- mmap alignment must be maintained: (p_offset % PAGE_SIZE) == (p_vaddr % PAGE_SIZE)
- Section headers must reflect the new offsets
"""

from dataclasses import dataclass
from pathlib import Path

from .types import (
    ProgramHeader,
    SectionHeader,
    PAGE_SIZE,
    PT_LOAD,
    SHT_NOBITS,
    round_up_to_page,
    round_down_to_page,
)
from .surgery import ElfSurgery, SectionInfo
from .phdr_manager import ProgramHeaderManager


@dataclass
class ZeroPageResult:
    """Result of zero-page optimization."""

    success: bool
    bytes_saved: int
    phdr_relocated: bool
    pages_zeroed: int
    error: str | None = None


def calculate_aligned_range(
    vaddr: int, size: int, page_size: int = PAGE_SIZE
) -> tuple[int, int]:
    """Calculate the page-aligned range within a section.

    Args:
        vaddr: Section virtual address
        size: Section size
        page_size: OS page size

    Returns:
        Tuple of (aligned_vaddr, aligned_size) - the range of full pages
        that can be zeroed. Returns (0, 0) if no full pages exist.
    """
    end = vaddr + size

    # Round up start to next page boundary
    aligned_start = round_up_to_page(vaddr, page_size)

    # Round down end to previous page boundary
    aligned_end = round_down_to_page(end, page_size)

    if aligned_start >= aligned_end:
        # No full pages to zero
        return (0, 0)

    return (aligned_start, aligned_end - aligned_start)


@dataclass
class SegmentPiece:
    """A piece of a split PT_LOAD segment."""

    vaddr: int
    memsz: int
    offset: int
    filesz: int

    @property
    def is_nobits(self) -> bool:
        """Check if this is a NOBITS (file-less) segment."""
        return self.filesz == 0 and self.memsz > 0


def split_load_segment(
    original: ProgramHeader,
    section_vaddr: int,
    section_size: int,
    aligned_vaddr: int,
    aligned_size: int,
    aligned_offset: int,
) -> list[SegmentPiece]:
    """Split a PT_LOAD segment around a zeroed region.

    The original segment is split into up to 4 pieces:
    1. Content before section (if any)
    2. Section prefix (if unaligned start)
    3. NOBITS region covering aligned portion AND any unaligned suffix
       (We extend NOBITS to cover suffix to avoid offset collision)
    4. Content after section (if any)

    Note: We do NOT create a separate suffix segment because that would
    cause a NOBITS collision - both the NOBITS and suffix segments would
    have the same file offset after content removal. Instead, we extend
    the NOBITS memsz to cover the suffix. The suffix content is still
    present in the file (at aligned_offset after removal), but it's not
    mapped by a separate PT_LOAD.

    Args:
        original: The PT_LOAD segment being split
        section_vaddr: Virtual address of the section being zeroed
        section_size: Size of the section
        aligned_vaddr: Virtual address of the aligned (zero-page) region
        aligned_size: Size of the aligned region
        aligned_offset: File offset of the aligned region (before removal)

    Returns:
        List of SegmentPieces representing the split
    """
    pieces: list[SegmentPiece] = []

    section_end_vaddr = section_vaddr + section_size
    section_offset = aligned_offset - (aligned_vaddr - section_vaddr)

    # Piece 1: Content before section (if any)
    if original.p_vaddr < section_vaddr:
        pre_vsize = section_vaddr - original.p_vaddr
        pre_fsize = min(pre_vsize, original.p_filesz)
        pieces.append(
            SegmentPiece(
                vaddr=original.p_vaddr,
                memsz=pre_vsize,
                offset=original.p_offset,
                filesz=pre_fsize,
            )
        )

    # Piece 2: Section prefix (if unaligned start)
    if section_vaddr < aligned_vaddr:
        prefix_size = aligned_vaddr - section_vaddr
        pieces.append(
            SegmentPiece(
                vaddr=section_vaddr,
                memsz=prefix_size,
                offset=section_offset,
                filesz=prefix_size,
            )
        )

    # Piece 3: NOBITS region covering aligned portion AND suffix
    # We extend memsz to cover to section end to avoid creating a separate
    # suffix segment that would have the same file offset (NOBITS collision)
    nobits_vsize = section_end_vaddr - aligned_vaddr
    pieces.append(
        SegmentPiece(
            vaddr=aligned_vaddr,
            memsz=nobits_vsize,  # Extended to cover suffix
            offset=aligned_offset,
            filesz=0,  # NOBITS - no file content
        )
    )

    # Piece 4: Content after section (if any)
    target_end = original.p_vaddr + original.p_memsz
    if section_end_vaddr < target_end:
        post_vaddr = section_end_vaddr
        post_vsize = target_end - post_vaddr
        # Calculate post-section file content
        section_file_end = section_offset + section_size
        post_offset = section_file_end - aligned_size  # Shifted down
        post_fsize = (original.p_offset + original.p_filesz) - section_file_end
        if post_fsize > 0:
            pieces.append(
                SegmentPiece(
                    vaddr=post_vaddr,
                    memsz=post_vsize,
                    offset=post_offset,
                    filesz=post_fsize,
                )
            )

    return pieces


def conservative_zero_page(
    surgery: ElfSurgery,
    section_name: str = ".hip_fatbin",
    spare_slots: int = 16,
) -> ZeroPageResult:
    """Apply conservative zero-page optimization to a section.

    This removes page-aligned content from the file while maintaining
    ELF validity. Only full pages are removed; partial pages are preserved.

    Args:
        surgery: ElfSurgery instance to operate on
        section_name: Section to zero-page (default: .hip_fatbin)
        spare_slots: Number of spare PHDR slots when relocating

    Returns:
        ZeroPageResult with operation outcome
    """
    # Find the target section
    section = surgery.find_section(section_name)
    if section is None:
        return ZeroPageResult(
            success=False,
            bytes_saved=0,
            phdr_relocated=False,
            pages_zeroed=0,
            error=f"Section '{section_name}' not found",
        )

    # Check if already zero-paged (NOBITS)
    if section.header.is_nobits:
        return ZeroPageResult(
            success=True,
            bytes_saved=0,
            phdr_relocated=False,
            pages_zeroed=0,
            error=f"Section '{section_name}' is already NOBITS (already zero-paged)",
        )

    page_size = surgery.page_size
    section_vaddr = section.header.sh_addr
    section_size = section.header.sh_size
    section_offset = section.header.sh_offset

    # Calculate aligned range
    aligned_vaddr, aligned_size = calculate_aligned_range(
        section_vaddr, section_size, page_size
    )

    if aligned_size == 0:
        # Section is too small or misaligned - no optimization possible
        return ZeroPageResult(
            success=True,
            bytes_saved=0,
            phdr_relocated=False,
            pages_zeroed=0,
            error="Section too small or misaligned - no full pages to zero",
        )

    # When the section ends exactly on a page boundary (suffix==0), removing
    # ALL aligned content causes a NOBITS collision: the NOBITS segment and
    # the post-section segment would share the same file offset. Fix: leave
    # one page of content behind so the post-section piece starts one page
    # higher. Cost: one fewer page of savings (4KB), but the ELF structure
    # is completely standard with no synthetic offsets.
    section_end = section_vaddr + section_size
    suffix = section_end - round_down_to_page(section_end, page_size)
    if suffix == 0:
        aligned_size -= page_size
        if aligned_size == 0:
            return ZeroPageResult(
                success=True,
                bytes_saved=0,
                phdr_relocated=False,
                pages_zeroed=0,
                error="Section page-aligned end with only one page - skipping to avoid NOBITS collision",
            )

    aligned_offset = section_offset + (aligned_vaddr - section_vaddr)
    pages_to_zero = aligned_size // page_size

    # Find the PT_LOAD containing the section
    target_result = surgery.find_phdr_containing_vaddr(section_vaddr)
    if target_result is None:
        return ZeroPageResult(
            success=False,
            bytes_saved=0,
            phdr_relocated=False,
            pages_zeroed=0,
            error="Could not find PT_LOAD containing section",
        )

    target_idx, target_load = target_result

    if target_load.p_type != PT_LOAD:
        return ZeroPageResult(
            success=False,
            bytes_saved=0,
            phdr_relocated=False,
            pages_zeroed=0,
            error=f"Segment {target_idx} is not PT_LOAD",
        )

    # Remove the aligned bytes from file
    old_size = len(surgery.data)
    del surgery.data[aligned_offset : aligned_offset + aligned_size]
    bytes_removed = old_size - len(surgery.data)

    # Adjust ELF header offsets for removed bytes.
    # This must happen before ProgramHeaderManager.apply() so it sees
    # correct e_phoff when calculating available space for in-place writes.
    ehdr = surgery.ehdr
    if ehdr.e_phoff > aligned_offset:
        surgery._ehdr.e_phoff -= aligned_size
    if ehdr.e_shoff > aligned_offset:
        surgery._ehdr.e_shoff -= aligned_size

    # Split the target PT_LOAD into pieces
    pieces = split_load_segment(
        target_load,
        section_vaddr,
        section_size,
        aligned_vaddr,
        aligned_size,
        aligned_offset,
    )

    # Build new program header list
    new_phdrs: list[ProgramHeader] = []

    for idx, phdr in surgery.iter_program_headers():
        if idx != target_idx:
            # Not the target - adjust offset if after removed bytes
            if phdr.p_offset > aligned_offset:
                new_phdr = ProgramHeader(
                    p_type=phdr.p_type,
                    p_flags=phdr.p_flags,
                    p_offset=phdr.p_offset - aligned_size,
                    p_vaddr=phdr.p_vaddr,
                    p_paddr=phdr.p_paddr,
                    p_filesz=phdr.p_filesz,
                    p_memsz=phdr.p_memsz,
                    p_align=phdr.p_align,
                )
                new_phdrs.append(new_phdr)
            else:
                new_phdrs.append(phdr)
        else:
            # Replace target with split pieces
            for piece in pieces:
                # Create PT_LOAD for this piece
                new_phdr = ProgramHeader(
                    p_type=PT_LOAD,
                    p_flags=target_load.p_flags,
                    p_offset=piece.offset,
                    p_vaddr=piece.vaddr,
                    p_paddr=piece.vaddr,
                    p_filesz=piece.filesz,
                    p_memsz=piece.memsz,
                    p_align=target_load.p_align,
                )
                new_phdrs.append(new_phdr)

    # Use ProgramHeaderManager to write the new headers
    # This handles overflow and relocation automatically
    surgery._phdrs = new_phdrs
    surgery._ehdr.e_phnum = len(new_phdrs)

    # Write all program headers
    manager = ProgramHeaderManager(surgery, spare_slots=spare_slots)
    manager._phdrs = new_phdrs
    phdr_result = manager.apply()

    # Update section headers
    # (e_phoff and e_shoff were already adjusted above, before apply())
    for sect in surgery.iter_sections():
        if sect.index == section.index:
            # Mark as NOBITS
            new_shdr = SectionHeader(
                sh_name=sect.header.sh_name,
                sh_type=SHT_NOBITS,
                sh_flags=sect.header.sh_flags,
                sh_addr=sect.header.sh_addr,
                sh_offset=section_offset,  # Keep original offset (irrelevant for NOBITS)
                sh_size=sect.header.sh_size,
                sh_link=sect.header.sh_link,
                sh_info=sect.header.sh_info,
                sh_addralign=sect.header.sh_addralign,
                sh_entsize=sect.header.sh_entsize,
            )
            surgery.update_section_header(sect.index, new_shdr)
        elif sect.header.sh_offset > aligned_offset:
            # Adjust offset
            new_shdr = SectionHeader(
                sh_name=sect.header.sh_name,
                sh_type=sect.header.sh_type,
                sh_flags=sect.header.sh_flags,
                sh_addr=sect.header.sh_addr,
                sh_offset=sect.header.sh_offset - aligned_size,
                sh_size=sect.header.sh_size,
                sh_link=sect.header.sh_link,
                sh_info=sect.header.sh_info,
                sh_addralign=sect.header.sh_addralign,
                sh_entsize=sect.header.sh_entsize,
            )
            surgery.update_section_header(sect.index, new_shdr)

    # Update ELF header
    surgery.update_elf_header()

    return ZeroPageResult(
        success=True,
        bytes_saved=bytes_removed,
        phdr_relocated=phdr_result.relocated,
        pages_zeroed=pages_to_zero,
    )


def zero_page_section(
    input_path: Path,
    output_path: Path,
    section_name: str = ".hip_fatbin",
    spare_slots: int = 16,
) -> ZeroPageResult:
    """Apply zero-page optimization to a file.

    Convenience function that loads, modifies, and saves an ELF binary.

    Args:
        input_path: Input ELF binary
        output_path: Output ELF binary
        section_name: Section to zero-page
        spare_slots: Number of spare PHDR slots

    Returns:
        ZeroPageResult with operation outcome
    """
    surgery = ElfSurgery.load(input_path)
    result = conservative_zero_page(surgery, section_name, spare_slots)

    if result.success:
        surgery.save_preserving_mode(output_path, input_path.stat().st_mode)

    return result