nutcracker.sh

#!/bin/bash
# ============================================================================
# NutCracker Compiler Build and Run Script
# ============================================================================

set -e  # Exit on error

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
MAGENTA='\033[0;35m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# Directories
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BUILD_DIR="$ROOT_DIR/build"
APPS_DIR="$ROOT_DIR/apps"
OUTPUT_DIR="$ROOT_DIR/vdsa_output"
P4C_DIR="$ROOT_DIR/third_party/p4mlir/third_party/p4c"
P4_INCLUDE_DIR="$ROOT_DIR/p4include"

# Tools
P4MLIR_TRANSLATE="$P4C_DIR/build/p4mlir-translate"
NUTCRACKER_OPT="$BUILD_DIR/bin/nutcracker-opt"
MAPPER="$ROOT_DIR/deps/mapper/target/release/mapper"
CLANG="${CLANG:-$ROOT_DIR/third_party/llvm/build/bin/clang}"

# ============================================================================
# Helper Functions
# ============================================================================

print_header() {
    local title="$1"
    echo ""
    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "🚀 ${BOLD}${GREEN}$title${NC}"
    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo ""
}

print_section() {
    echo ""
    echo -e "📦 ${BOLD}${CYAN}$1${NC}"
    echo ""
}

print_step() {
    echo -e "${GREEN}  ▸${NC} $1"
}

print_substep() {
    echo -e "    ${CYAN}→${NC} $1"
}

print_error() {
    echo -e "${RED}❌ Error:${NC} $1" >&2
}

print_success() {
    echo -e "${GREEN}✅${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}⚠️${NC}  $1"
}

print_info() {
    echo -e "${BLUE}ℹ️${NC}  $1"
}

print_file() {
    echo -e "    ${MAGENTA}📄${NC} $1"
}

print_dir() {
    echo -e "    ${BLUE}📁${NC} $1"
}

# ============================================================================
# Setup Functions
# ============================================================================

setup_p4_includes() {
    print_header "Setting up P4 Include Directory"
    
    mkdir -p "$P4_INCLUDE_DIR"
    
    # Check if nc.p4 exists
    if [ ! -f "$P4_INCLUDE_DIR/nc.p4" ]; then
        print_warning "nc.p4 not found in p4include/"
        print_step "Creating nc.p4 architecture definition..."
        
        cat > "$P4_INCLUDE_DIR/nc.p4" << 'EOF'
// nc.p4 - NutCracker Architecture Definition
#ifndef _NC_P4_
#define _NC_P4_

#include <core.p4>

// ============================================================================
// Standard Metadata
// ============================================================================

struct nc_standard_metadata_t {
    bit<9>  ingress_port;
    bit<9>  egress_port;
    bit<32> packet_length;
    bit<1>  drop;
}

// ============================================================================
// Extern Functions (Fixed-Function Accelerators)
// ============================================================================

extern Counter<I> {
    Counter(bit<32> size);
    void count(in I index);
}

extern Hash5Tuple {
    Hash5Tuple();
    void apply(out bit<32>, in bit<32>, in bit<32>, 
               in bit<16>, in bit<16>, in bit<8>);
}

// ============================================================================
// Programmable Block Type Definitions
// ============================================================================

parser MainParserT<MH, MM>(
    packet_in pkt,
    out MH main_hdr,
    inout MM main_meta,
    in nc_standard_metadata_t standard_meta
);

control MainControlT<MH, MM>(
    inout MH main_hdr,
    inout MM main_meta,
    inout nc_standard_metadata_t standard_meta
);

control MainDeparserT<MH, MM>(
    packet_out pkt,
    in MH main_hdr,
    in MM main_meta,
    in nc_standard_metadata_t standard_meta
);

// ============================================================================
// NC Pipeline Package
// ============================================================================

package NC_PIPELINE<MH, MM>(
    MainParserT<MH, MM> main_parser,
    MainControlT<MH, MM> main_control,
    MainDeparserT<MH, MM> main_deparser
);

#endif  // _NC_P4_
EOF
        print_success "Created nc.p4"
    else
        print_success "nc.p4 already exists"
    fi
}

# ============================================================================
# Build Functions
# ============================================================================

build_llvm() {
    print_header "Building LLVM"
    
    if [ -d "$ROOT_DIR/third_party/llvm/build" ]; then
        print_warning "LLVM already built, skipping..."
        return 0
    fi
    
    print_step "Running LLVM build script..."
    cd "$ROOT_DIR"
    bash build_scripts/build_llvm.sh
    print_success "LLVM built successfully"
}

build_p4mlir() {
    print_header "Building P4MLIR"
    
    if [ -f "$P4MLIR_TRANSLATE" ]; then
        print_warning "P4MLIR already built, skipping..."
        return 0
    fi
    
    print_step "Building P4MLIR..."
    cd "$P4C_DIR"
    mkdir -p build
    cd build
    cmake .. -DENABLE_MLIR=ON
    make -j$(nproc)
    print_success "P4MLIR built successfully"
}

build_nutcracker() {
    print_header "Building NutCracker"
    
    print_step "Configuring CMake..."
    mkdir -p "$BUILD_DIR"
    cd "$BUILD_DIR"
    
    cmake .. -DCMAKE_BUILD_TYPE=Debug 
    
    print_step "Building NutCracker..."
    make -j$(nproc)
    
    print_success "NutCracker built successfully"
}

# ============================================================================
# Compilation Functions
# ============================================================================

compile_p4_to_mlir() {
    local p4_file=$1
    local output_file=$2
    local force=${3:-false}
    
    # Skip if MLIR exists and not forcing
    if [ -f "$output_file" ] && [ "$force" = false ]; then
        print_substep "MLIR already exists, skipping: $(basename $output_file)"
        return 0
    fi
    
    print_step "P4 → P4HIR MLIR: $(basename $p4_file)"
    print_substep "Input:  $p4_file"
    print_substep "Output: $output_file"
    
    if [ ! -f "$P4MLIR_TRANSLATE" ]; then
        print_error "P4MLIR translator not found. Please build P4MLIR first."
        exit 1
    fi
    
    # Build P4 include paths
    local include_args=""
    
    if [ -d "$P4_INCLUDE_DIR" ]; then
        include_args="-I$P4_INCLUDE_DIR"
    fi
    
    local p4c_include="$P4C_DIR/p4include"
    if [ -d "$p4c_include" ]; then
        include_args="$include_args -I$p4c_include"
    fi
    
    local app_dir=$(dirname "$p4_file")
    if [ -d "$app_dir" ]; then
        include_args="$include_args -I$app_dir"
    fi
    
    "$P4MLIR_TRANSLATE" $include_args --typeinference-only "$p4_file" > "$output_file" 2>&1
    
    if [ $? -eq 0 ]; then
        local line_count=$(wc -l < "$output_file")
        local op_count=$(grep -c "p4hir\." "$output_file" || echo "0")
        print_success "Generated P4HIR MLIR ($line_count lines, ~$op_count ops)"
    else
        print_error "P4 to MLIR compilation failed"
        cat "$output_file"
        exit 1
    fi
}

run_partition_pass() {
    local mlir_file=$1
    local output_dir=$2  # Will be vdsa_output/
    
    print_step "Step 2: Partitioning P4HIR into Basic Blocks"
    print_substep "Input:  $mlir_file"
    print_substep "Output: $output_dir"
    
    if [ ! -f "$NUTCRACKER_OPT" ]; then
        print_error "nutcracker-opt not found. Please build NutCracker first."
        exit 1
    fi
    
    if [ ! -f "$mlir_file" ]; then
        print_error "Input MLIR file not found: $mlir_file"
        exit 1
    fi
    
    # Get directories
    local mlir_dir=$(dirname "$mlir_file")
    local mlir_name=$(basename "$mlir_file")
    
    # Change to MLIR directory and run pass
    cd "$mlir_dir"
    
    print_substep "Running partition pass (creates vdsa_output in current dir)..."
    
    # Pass creates vdsa_output/ in current directory
    "$NUTCRACKER_OPT" --p4hir-partition "$mlir_name" 2>&1
    local result=$?
    
    cd "$ROOT_DIR"
    
    if [ $result -ne 0 ]; then
        print_error "Partition pass failed"
        return 1
    fi
    
    # Pass created vdsa_output in MLIR directory
    local pass_output_dir="$mlir_dir/vdsa_output"
    
    if [ ! -d "$pass_output_dir" ]; then
        print_error "Pass did not create vdsa_output directory"
        return 1
    fi
    
    # Move to ROOT_DIR/vdsa_output/
    rm -rf "$output_dir"
    mv "$pass_output_dir" "$output_dir"
    
    if [ ! -d "$output_dir" ]; then
        print_error "Failed to move output"
        return 1
    fi
    
    # Count blocks
    local block_count=$(find "$output_dir" -maxdepth 1 -type d -name "block*" 2>/dev/null | wc -l)
    
    if [ "$block_count" -eq 0 ]; then
        print_error "No blocks generated"
        return 1
    fi
    
    print_success "Partitioned into $block_count basic blocks"
    echo ""
}

# run_vffs_coarse_grained removed: P4HIRToVDRMT/VDPP now embed the
# coarse-grained finalization directly.

run_fine_grained_lowering() {
    local mlir_file=$1
    local output_dir=$2

    print_step "Step 4a: Fine-Grained Lowering (vDRMT → bf3drmt)"
    print_substep "Partition dir: $output_dir"

    if [ ! -f "$NUTCRACKER_OPT" ]; then
        print_error "nutcracker-opt not found."
        exit 1
    fi

    if [ ! -d "$output_dir" ]; then
        print_error "Partition output directory not found: $output_dir"
        return 1
    fi

    cd "$ROOT_DIR"

    # Export NUTCRACKER_ROOT so the pass can find the egglog template
    export NUTCRACKER_ROOT="$ROOT_DIR"

    print_substep "Running vDRMT → bf3drmt pass..."
    "$NUTCRACKER_OPT" --vdrmt-to-bf3drmt "$mlir_file" 2>&1
    local result=$?

    if [ $result -ne 0 ]; then
        print_warning "vDRMT fine-grained lowering reported errors (continuing)"
    fi

    local bf3drmt_count=$(find "$output_dir" -name "bf3drmt.mlir" 2>/dev/null | wc -l)

    if [ "$bf3drmt_count" -gt 0 ]; then
        echo -e "    ${GREEN}✅ bf3drmt:${NC} $bf3drmt_count blocks"
        for block_dir in "$output_dir"/block*; do
            if [ -f "$block_dir/bf3drmt.mlir" ]; then
                echo -e "      ${GREEN}→${NC} $(basename $block_dir)/bf3drmt.mlir"
            fi
        done
    else
        echo -e "    ${YELLOW}⚠️  bf3drmt:${NC} 0 blocks"
    fi

    echo ""
    return 0
}

run_fine_grained_lowering_vdpp() {
    local mlir_file=$1
    local output_dir=$2

    print_step "Step 4b: Fine-Grained Lowering (vDPP → bf3dpa)"
    print_substep "Partition dir: $output_dir"

    if [ ! -f "$NUTCRACKER_OPT" ]; then
        print_error "nutcracker-opt not found."
        exit 1
    fi

    if [ ! -d "$output_dir" ]; then
        print_error "Partition output directory not found: $output_dir"
        return 1
    fi

    # Check if any vdpp.mlir blocks exist
    local vdpp_count=$(find "$output_dir" -name "vdpp.mlir" 2>/dev/null | wc -l)
    if [ "$vdpp_count" -eq 0 ]; then
        print_substep "No vDPP blocks found, skipping"
        echo ""
        return 0
    fi

    cd "$ROOT_DIR"
    export NUTCRACKER_ROOT="$ROOT_DIR"

    print_substep "Running vDPP → bf3dpa pass..."
    "$NUTCRACKER_OPT" --vdpp-to-bf3dpa "$mlir_file" 2>&1
    local result=$?

    if [ $result -ne 0 ]; then
        print_warning "vDPP fine-grained lowering reported errors (continuing)"
    fi

    local bf3dpa_count=$(find "$output_dir" -name "bf3dpa.mlir" 2>/dev/null | wc -l)

    if [ "$bf3dpa_count" -gt 0 ]; then
        echo -e "    ${CYAN}✅ bf3dpa:${NC} $bf3dpa_count blocks"
        for block_dir in "$output_dir"/block*; do
            if [ -f "$block_dir/bf3dpa.mlir" ]; then
                echo -e "      ${CYAN}→${NC} $(basename $block_dir)/bf3dpa.mlir"
            fi
        done
    else
        echo -e "    ${YELLOW}⚠️  bf3dpa:${NC} 0 blocks"
    fi

    [ -f "$output_dir/dpa_handler.c" ] && \
        echo -e "    ${CYAN}✅ dpa_handler.c generated${NC}"
    [ -f "$output_dir/arm_handler.c" ] && \
        echo -e "    ${CYAN}✅ arm_handler.c generated${NC}"

    echo ""
    return 0
}

run_lowering_passes() {
    local mlir_file=$1
    local output_dir=$2
    
    print_step "Step 3: Lowering P4HIR to Target Dialects"
    print_substep "Partition dir: $output_dir"
    
    if [ ! -f "$NUTCRACKER_OPT" ]; then
        print_error "nutcracker-opt not found."
        exit 1
    fi
    
    if [ ! -d "$output_dir" ]; then
        print_error "Partition output directory not found: $output_dir"
        return 1
    fi
    
    cd "$ROOT_DIR"
    
    # Run vDRMT lowering (handles simple blocks)
    print_substep "Trying vDRMT lowering (Match-Action)..."
    "$NUTCRACKER_OPT" --p4hir-to-vdrmt "$mlir_file" 2>&1
    local vdrmt_result=$?
    
    # Run vDPP lowering (handles complex blocks)
    print_substep "Trying vDPP lowering (Programmable)..."
    "$NUTCRACKER_OPT" --p4hir-to-vdpp "$mlir_file" 2>&1
    local vdpp_result=$?
    
    # Check results
    local vdrmt_count=$(find "$output_dir" -name "vdrmt.mlir" 2>/dev/null | wc -l)
    local vdpp_count=$(find "$output_dir" -name "vdpp.mlir" 2>/dev/null | wc -l)
    
    echo ""
    print_substep "Lowering Results:"
    
    if [ "$vdrmt_count" -gt 0 ]; then
        echo -e "    ${GREEN}✅ vDRMT:${NC} $vdrmt_count blocks"
        for block_dir in "$output_dir"/block*; do
            if [ -f "$block_dir/vdrmt.mlir" ]; then
                echo -e "      ${GREEN}→${NC} $(basename $block_dir)/vdrmt.mlir"
            fi
        done
    else
        echo -e "    ${YELLOW}⚠️  vDRMT:${NC} 0 blocks"
    fi
    
    if [ "$vdpp_count" -gt 0 ]; then
        echo -e "    ${BLUE}✅ vDPP:${NC}  $vdpp_count blocks"
        for block_dir in "$output_dir"/block*; do
            if [ -f "$block_dir/vdpp.mlir" ]; then
                echo -e "      ${BLUE}→${NC} $(basename $block_dir)/vdpp.mlir"
            fi
        done
    else
        echo -e "    ${YELLOW}⚠️  vDPP:${NC}  0 blocks"
    fi
    
    local total_lowered=$((vdrmt_count + vdpp_count))
    
    if [ "$total_lowered" -eq 0 ]; then
        print_warning "No blocks were lowered to any target"
        return 1
    fi
    
    print_success "Lowered $total_lowered blocks total"
    echo ""
    
    return 0
}

run_emit_handler_code() {
    local mlir_file=$1
    local output_dir=$2

    print_step "Step 5.5: Emit Handler Code (mapping-aware)"

    if [ ! -f "$NUTCRACKER_OPT" ]; then
        print_error "nutcracker-opt not found."
        return 1
    fi

    cd "$ROOT_DIR"
    export NUTCRACKER_ROOT="$ROOT_DIR"

    "$NUTCRACKER_OPT" --emit-handler-code "$mlir_file" 2>&1
    local result=$?

    if [ $result -ne 0 ]; then
        print_warning "emit-handler-code reported errors (continuing)"
    fi

    [ -f "$output_dir/nc_types.h"    ] && echo -e "    ${CYAN}✅ nc_types.h${NC}"
    [ -f "$output_dir/dpa_handler.c" ] && echo -e "    ${CYAN}✅ dpa_handler.c${NC}"
    [ -f "$output_dir/arm_handler.c" ] && echo -e "    ${CYAN}✅ arm_handler.c${NC}"
    [ -f "$output_dir/Makefile"      ] && echo -e "    ${CYAN}✅ Makefile${NC}"
    echo ""
    return 0
}

run_mapper() {
    local output_dir=$1
    local mode=${2:-balanced}   # balanced | min-latency | max-throughput | min-energy
    local mapper_egg="$output_dir/mapper.egg"

    print_step "Step 5: Deployment Mapping (mapper)"
    print_substep "Input:  $mapper_egg"
    print_substep "Mode:   $mode"

    if [ ! -f "$MAPPER" ]; then
        print_warning "mapper binary not found: $MAPPER"
        print_info  "Build it with: cd deps/mapper && cargo build --release"
        return 0
    fi

    if [ ! -f "$mapper_egg" ]; then
        print_warning "mapper.egg not found (fine-grained lowering may have failed)"
        return 0
    fi

    local mapper_out="$output_dir/mapping.txt"
    "$MAPPER" "$mapper_egg" --mode "$mode" --per-objective > "$mapper_out" 2>&1
    local result=$?

    if [ $result -ne 0 ]; then
        print_error "Mapper failed (exit $result)"
        cat "$mapper_out" >&2
        return 1
    fi

    # Print the mapping result inline
    cat "$mapper_out"

    print_success "Mapping written to $mapper_out"
    echo ""
    return 0
}

generate_visualization() {
    local dot_file=$1
    local output_pdf=$2

    if [ ! -f "$dot_file" ]; then
        print_warning "DOT file not found: $dot_file"
        return 1
    fi

    print_step "Step 6: Generating Control Flow Visualization"
    
    if command -v dot &> /dev/null; then
        dot -Tpdf "$dot_file" -o "$output_pdf" 2>/dev/null
        print_success "Generated PDF: $(basename $output_pdf)"
        
        dot -Tpng "$dot_file" -o "${output_pdf%.pdf}.png" 2>/dev/null
        print_success "Generated PNG: $(basename ${output_pdf%.pdf}.png)"
    else
        print_warning "Graphviz not installed. Skipping visualization."
        print_info "Install with: sudo apt-get install graphviz"
    fi
    echo ""
}

# ============================================================================
# C App Functions (C → LLVM IR → vDRMT/vDPP)
# ============================================================================

precompile_c_app() {
    local app_name=$1
    local app_dir="$APPS_DIR/$app_name"
    local c_file="$app_dir/${app_name}.c"
    local ll_file="$app_dir/arm.ll"
    local force=${2:-false}

    if [ -f "$ll_file" ] && [ "$force" = false ]; then
        print_info "LLVM IR already exists for $app_name (use --force to regenerate)"
        return 0
    fi

    print_section "Precompiling C app: $app_name"
    print_step "C → LLVM IR: $(basename $c_file)"
    print_substep "Clang: $CLANG"
    print_substep "Output: $ll_file"

    if [ ! -f "$CLANG" ]; then
        print_error "clang not found: $CLANG"
        exit 1
    fi

    local ll_raw="$app_dir/arm_raw.ll"
    local OPT="${OPT:-$ROOT_DIR/third_party/llvm/build/bin/opt}"

    "$CLANG" -O0 -S -emit-llvm \
        -fno-discard-value-names \
        -o "$ll_raw" "$c_file" 2>&1

    if [ $? -ne 0 ]; then
        print_error "C compilation failed"
        return 1
    fi

    # Promote alloca-based variables to SSA (simplifies provenance tracking).
    # Strip 'optnone' first so mem2reg can run (it's added by -O0).
    print_substep "Running mem2reg pass..."
    sed 's/ optnone//g; s/noinline //g' "$ll_raw" | \
        "$OPT" -passes=mem2reg -S -o "$ll_file" 2>&1

    if [ $? -ne 0 ]; then
        print_warning "mem2reg failed, using raw IR"
        cp "$ll_raw" "$ll_file"
    fi

    local line_count=$(wc -l < "$ll_file")
    print_success "Generated LLVM IR ($line_count lines)"
    echo ""
}

compile_c_app() {
    local app_name=$1
    local app_dir="$APPS_DIR/$app_name"
    local ll_file="$app_dir/arm.ll"
    local output_dir="$OUTPUT_DIR"

    print_header "Compiling C app: $app_name"

    if [ ! -f "$ll_file" ]; then
        print_error "LLVM IR not found: $ll_file (run precompile first)"
        exit 1
    fi

    if [ ! -f "$NUTCRACKER_OPT" ]; then
        print_error "nutcracker-opt not found. Please build NutCracker first."
        exit 1
    fi

    rm -rf "$output_dir"

    # Run vDPP lowering (non-fatal: some ops may be unhandled)
    print_step "Step 1: LLVM IR → vDPP"
    echo "" | "$NUTCRACKER_OPT" \
        "--llvm-ir-to-vdpp=input-ll=$ll_file output-dir=$output_dir" \
        /dev/null 2>&1 || true
    local vdpp_result=$?

    # Run vDRMT lowering
    print_step "Step 2: LLVM IR → vDRMT"
    # Run per block since each block has its own arm.ll from vDPP
    local block_idx=0
    for block_dir in "$output_dir"/block*; do
        if [ -f "$block_dir/arm.ll" ]; then
            echo "" | "$NUTCRACKER_OPT" \
                "--llvm-ir-to-vdrmt=input-ll=$block_dir/arm.ll output-dir=$block_dir/.vdrmt_tmp" \
                /dev/null 2>&1
            # Move the generated vdrmt.mlir up from .vdrmt_tmp/block0/
            if [ -f "$block_dir/.vdrmt_tmp/block0/vdrmt.mlir" ]; then
                mv "$block_dir/.vdrmt_tmp/block0/vdrmt.mlir" "$block_dir/vdrmt.mlir"
            fi
            rm -rf "$block_dir/.vdrmt_tmp"
            block_idx=$((block_idx + 1))
        fi
    done

    # If vDPP didn't create arm.ll blocks (single-function C), run vDRMT on original
    if [ "$block_idx" -eq 0 ]; then
        print_substep "No per-block arm.ll found; running vDRMT on original LLVM IR"
        mkdir -p "$output_dir/block0"
        echo "" | "$NUTCRACKER_OPT" \
            "--llvm-ir-to-vdrmt=input-ll=$ll_file output-dir=$output_dir/.vdrmt_tmp" \
            /dev/null 2>&1
        if [ -d "$output_dir/.vdrmt_tmp" ]; then
            for b in "$output_dir/.vdrmt_tmp"/block*; do
                local bname=$(basename "$b")
                mkdir -p "$output_dir/$bname"
                [ -f "$b/vdrmt.mlir" ] && mv "$b/vdrmt.mlir" "$output_dir/$bname/vdrmt.mlir"
            done
            rm -rf "$output_dir/.vdrmt_tmp"
        fi
    fi

    # Summary
    local vdrmt_count=$(find "$output_dir" -name "vdrmt.mlir" 2>/dev/null | wc -l)
    local vdpp_count=$(find "$output_dir" -name "vdpp.mlir" 2>/dev/null | wc -l)
    local block_count=$(find "$output_dir" -maxdepth 1 -type d -name "block*" 2>/dev/null | wc -l)

    print_header "Compilation Summary"
    echo -e "${BOLD}📦 Application:${NC} $app_name"
    echo -e "${BOLD}🧩 Blocks:${NC}      $block_count"
    echo -e "${BOLD}📂 Output:${NC}      $output_dir"
    echo ""
    echo -e "  ${GREEN}vDRMT blocks:${NC} $vdrmt_count"
    echo -e "  ${BLUE}vDPP blocks:${NC}  $vdpp_count"
    echo ""
    for block_dir in "$output_dir"/block*; do
        if [ -d "$block_dir" ]; then
            local bname=$(basename "$block_dir")
            echo -e "  ${CYAN}📦 $bname:${NC}"
            [ -f "$block_dir/vdpp.mlir"  ] && echo -e "    ${BLUE}📄 vdpp.mlir${NC}"
            [ -f "$block_dir/vdrmt.mlir" ] && echo -e "    ${GREEN}📄 vdrmt.mlir${NC}"
            [ -f "$block_dir/arm.ll"     ] && echo -e "    ${MAGENTA}📄 arm.ll${NC}"
        fi
    done
    echo ""
    print_success "Compilation complete!"
    echo ""
}

# ============================================================================
# Precompile Functions
# ============================================================================

precompile_app() {
    local app_name=$1
    local app_dir="$APPS_DIR/$app_name"
    local p4_file="$app_dir/${app_name}.p4"
    local c_file="$app_dir/${app_name}.c"
    local mlir_file="$app_dir/${app_name}.mlir"
    local force=${2:-false}

    # Dispatch to C path if no P4 but C exists
    if [ ! -f "$p4_file" ] && [ -f "$c_file" ]; then
        precompile_c_app "$app_name" "$force"
        return $?
    fi

    if [ ! -f "$p4_file" ]; then
        print_error "P4 file not found: $p4_file"
        return 1
    fi
    
    if [ -f "$mlir_file" ] && [ "$force" = false ]; then
        print_info "MLIR already exists for $app_name (use --force to regenerate)"
        return 0
    fi
    
    print_section "Precompiling: $app_name"
    compile_p4_to_mlir "$p4_file" "$mlir_file" "$force"
    echo ""
}

precompile_all_apps() {
    print_header "Precompiling All Applications (P4 → MLIR)"
    
    local force=false
    if [ "$1" = "--force" ]; then
        force=true
        print_info "Force regeneration enabled"
    fi
    
    setup_p4_includes
    
    local success_count=0
    local skip_count=0
    local fail_count=0
    
    for app_dir in "$APPS_DIR"/*; do
        if [ -d "$app_dir" ]; then
            local app_name=$(basename "$app_dir")
            local p4_file="$app_dir/${app_name}.p4"
            local mlir_file="$app_dir/${app_name}.mlir"
            
            if [ -f "$p4_file" ]; then
                if [ -f "$mlir_file" ] && [ "$force" = false ]; then
                    print_info "✓ $app_name (already compiled)"
                    ((skip_count++))
                else
                    if precompile_app "$app_name" "$force"; then
                        ((success_count++))
                    else
                        ((fail_count++))
                        print_error "Failed to precompile $app_name"
                    fi
                fi
            fi
        fi
    done
    
    print_header "Precompile Summary"
    echo -e "  ${GREEN}✅ Compiled:${NC} $success_count"
    echo -e "  ${BLUE}⏭️  Skipped:${NC}  $skip_count"
    echo -e "  ${RED}❌ Failed:${NC}   $fail_count"
    echo ""
    
    if [ $success_count -gt 0 ]; then
        print_success "Precompile complete! MLIR files ready for partition/lowering."
    fi
}

# ============================================================================
# Compile Application
# ============================================================================

compile_app() {
    local app_name=$1
    local app_dir="$APPS_DIR/$app_name"
    local p4_file="$app_dir/${app_name}.p4"
    local c_file="$app_dir/${app_name}.c"
    local mlir_file="$app_dir/${app_name}.mlir"
    local output_dir="$OUTPUT_DIR"

    # Dispatch to C path if no P4 but C exists
    if [ ! -f "$p4_file" ] && [ -f "$c_file" ]; then
        compile_c_app "$app_name"
        return $?
    fi

    print_header "Compiling: $app_name"

    if [ ! -f "$p4_file" ]; then
        print_error "P4 file not found: $p4_file"
        exit 1
    fi
    
    setup_p4_includes
    
    # Step 1: P4 → MLIR (skip if exists)
    compile_p4_to_mlir "$p4_file" "$mlir_file" false
    
    rm -rf "$output_dir"
    mkdir -p "$output_dir"
    
    # Step 2: Partition into blocks
    run_partition_pass "$mlir_file" "$output_dir"
    
    # Step 3: Lower to both vDRMT and vDPP
    # P4HIRToVDRMT/VDPP now embed the coarse-grained finalization (VFFA
    # instance lifting + helper-attr stripping); no separate step needed.
    run_lowering_passes "$mlir_file" "$output_dir"

    # Step 4: Fine-grained lowering (target-specific)
    if [ "$TARGET_BACKEND" = "bf2" ]; then
        # BF2: vDRMT → bf2drmt (DOCA 2.2)
        print_step "Step 4a: Fine-Grained Lowering (vDRMT → bf2drmt)"
        DOCA_TEMPLATE="$ROOT_DIR/templates/doca-2.2.toml" \
            "$NUTCRACKER_OPT" --vdrmt-to-bf2drmt 2>&1
    else
        # BF3 (default): vDRMT → bf3drmt + vDPP → bf3dpa
        run_fine_grained_lowering "$mlir_file" "$output_dir"
        run_fine_grained_lowering_vdpp "$mlir_file" "$output_dir"
    fi

    # Step 5: Deployment mapping
    run_mapper "$output_dir" "${MAPPER_MODE:-balanced}"

    # Step 5.5: Emit handler C code (mapping-aware) — BF3 only
    # BF2's doca_flow_pipeline.c is already generated by --vdrmt-to-bf2drmt.
    if [ "${TARGET_BACKEND:-bf3}" != "bf2" ]; then
        run_emit_handler_code "$mlir_file" "$output_dir"
    fi

    # Step 6: Visualize
    generate_visualization "$output_dir/dependency_graph.dot" "$output_dir/graph.pdf"
    
    print_header "Compilation Summary"
    
    if [ -d "$output_dir" ]; then
        local block_count=$(find "$output_dir" -maxdepth 1 -type d -name "block*" | wc -l)
        
        echo -e "${BOLD}📦 Application:${NC} $app_name"
        echo -e "${BOLD}🧩 Blocks:${NC}      $block_count"
        echo -e "${BOLD}📂 Output:${NC}      $output_dir"
        echo ""
        
        print_section "Block Statistics"
        local block_id=0
        local vdrmt_count=0
        local vdpp_count=0
        local bf3drmt_count=0
        local unmapped_count=0

        for block_dir in "$output_dir"/block*; do
            if [ -f "$block_dir/metadata.json" ]; then
                local ops=$(jq -r '.operations' "$block_dir/metadata.json" 2>/dev/null || echo "?")
                local has_if=$(jq -r '.hasConditionalBranch' "$block_dir/metadata.json" 2>/dev/null || echo "false")
                local has_table=$(jq -r '.hasTableApply' "$block_dir/metadata.json" 2>/dev/null || echo "false")

                echo -e "  ${CYAN}📦 Block $block_id:${NC} $ops ops"
                [ "$has_if" = "true" ] && echo -e "    ${YELLOW}🔀${NC} Conditional branch"
                [ "$has_table" = "true" ] && echo -e "    ${YELLOW}📊${NC} Table apply"

                local has_vdrmt=false
                local has_vdpp=false

                if [ -f "$block_dir/vdrmt.mlir" ]; then
                    echo -e "    ${GREEN}✅ vDRMT (Match-Action)${NC}"
                    echo -e "    ${GREEN}📄 vdrmt.mlir${NC}"
                    has_vdrmt=true
                    vdrmt_count=$((vdrmt_count + 1))
                fi

                if [ -f "$block_dir/bf3drmt.mlir" ]; then
                    echo -e "    ${MAGENTA}✅ bf3drmt (BF3 NIC ASIC)${NC}"
                    echo -e "    ${MAGENTA}📄 bf3drmt.mlir${NC}"
                    bf3drmt_count=$((bf3drmt_count + 1))
                fi

                if [ -f "$block_dir/vdpp.mlir" ]; then
                    echo -e "    ${BLUE}✅ vDPP (Programmable)${NC}"
                    echo -e "    ${BLUE}📄 vdpp.mlir${NC}"
                    has_vdpp=true
                    vdpp_count=$((vdpp_count + 1))
                fi

                if [ -f "$block_dir/bf3dpa.mlir" ]; then
                    echo -e "    ${CYAN}✅ bf3dpa (DPA/ARM)${NC}"
                    echo -e "    ${CYAN}📄 bf3dpa.mlir${NC}"
                fi

                if [ "$has_vdrmt" = false ] && [ "$has_vdpp" = false ]; then
                    echo -e "    ${RED}❌ Not lowered to any target${NC}"
                    unmapped_count=$((unmapped_count + 1))
                fi

                echo ""
                block_id=$((block_id + 1))
            fi
        done

        echo -e "${BOLD}Lowering Summary:${NC}"
        echo -e "  ${GREEN}vDRMT blocks:${NC}    $vdrmt_count"
        echo -e "  ${MAGENTA}bf3drmt blocks:${NC}  $bf3drmt_count"
        echo -e "  ${BLUE}vDPP blocks:${NC}     $vdpp_count"
        echo -e "  ${RED}Unmapped:${NC}        $unmapped_count"
        echo -e "  ${CYAN}Total blocks:${NC}    $block_id"
        echo ""
        
        print_section "Generated Files"
        echo -e "${BLUE}📁${NC} $output_dir/"
        echo -e "  ${MAGENTA}📄${NC} dependency_graph.dot"
        [ -f "$output_dir/graph.pdf" ]   && echo -e "  ${MAGENTA}📄${NC} graph.pdf"
        [ -f "$output_dir/graph.png" ]   && echo -e "  ${MAGENTA}📄${NC} graph.png"
        [ -f "$output_dir/mapper.egg" ]           && echo -e "  ${CYAN}📄${NC} mapper.egg"
        [ -f "$output_dir/mapping.txt" ]          && echo -e "  ${CYAN}📄${NC} mapping.txt"
        [ -f "$output_dir/doca_flow_pipeline.c" ] && echo -e "  ${CYAN}📄${NC} doca_flow_pipeline.c"
        [ -f "$output_dir/dpa_handler.c" ]       && echo -e "  ${CYAN}📄${NC} dpa_handler.c"
        [ -f "$output_dir/arm_handler.c" ]       && echo -e "  ${CYAN}📄${NC} arm_handler.c"

        for block_dir in "$output_dir"/block*; do
            if [ -d "$block_dir" ]; then
                local block_name=$(basename "$block_dir")
                echo -e "  ${BLUE}📁${NC} $block_name/"
                echo -e "    ${MAGENTA}📄${NC} p4hir.mlir"
                [ -f "$block_dir/vdrmt.mlir" ]    && echo -e "    ${GREEN}📄${NC} vdrmt.mlir"
                [ -f "$block_dir/bf3drmt.mlir" ]  && echo -e "    ${MAGENTA}📄${NC} bf3drmt.mlir"
                [ -f "$block_dir/vdpp.mlir" ]     && echo -e "    ${BLUE}📄${NC} vdpp.mlir"
                [ -f "$block_dir/bf3dpa.mlir" ]   && echo -e "    ${CYAN}📄${NC} bf3dpa.mlir"
                echo -e "    ${MAGENTA}📄${NC} metadata.json"
            fi
        done
        echo ""
    fi
    
    print_success "Compilation complete!"
    echo ""
}

# ============================================================================
# Clean Functions
# ============================================================================

clean_build() {
    print_header "Cleaning Build"
    
    print_step "Removing build directory..."
    rm -rf "$BUILD_DIR"
    
    print_step "Removing output directory..."
    rm -rf "$OUTPUT_DIR"
    
    print_success "Clean complete"
}

clean_mlir() {
    print_header "Cleaning MLIR Files"
    
    print_step "Removing generated MLIR files..."
    find "$APPS_DIR" -name "*.mlir" -type f -delete
    
    local count=$(find "$APPS_DIR" -name "*.mlir" -type f 2>/dev/null | wc -l)
    print_success "Removed MLIR files"
}

clean_all() {
    print_header "Deep Clean"
    
    clean_build
    clean_mlir
    
    print_step "Removing LLVM build..."
    rm -rf "$ROOT_DIR/third_party/llvm/build"
    
    print_step "Removing P4MLIR build..."
    rm -rf "$P4C_DIR/build"
    
    print_success "Deep clean complete"
}

# ============================================================================
# Main Menu
# ============================================================================

show_usage() {
    cat << EOF

${BOLD}${BLUE}🌰 NutCracker Compiler${NC} - Multi-Target vDSA Code Generator

${BOLD}USAGE:${NC}
    $0 <command> [options]

${BOLD}COMMANDS:${NC}
    ${GREEN}🔨 build${NC}              Build all components
    ${GREEN}🔨 build-llvm${NC}         Build LLVM only
    ${GREEN}🔨 build-p4mlir${NC}       Build P4MLIR only
    ${GREEN}🔨 build-nutcracker${NC}   Build NutCracker only
    
    ${GREEN}⚙️  setup${NC}              Setup P4 include directory
    
    ${GREEN}🎯 precompile${NC} <app>   Precompile P4 → MLIR for one app
    ${GREEN}🎯 precompile-all${NC}     Precompile all apps (P4 → MLIR)
    
    ${GREEN}🚀 compile${NC} <app>      Compile app (MLIR → vDRMT/vDPP)
    ${GREEN}🚀 compile-all${NC}        Compile all applications
    
    ${GREEN}🧹 clean${NC}              Clean build artifacts
    ${GREEN}🧹 clean-mlir${NC}         Clean generated MLIR files
    ${GREEN}🧹 clean-all${NC}          Deep clean (including dependencies)
    
    ${GREEN}▶️  run${NC} <app>          Build and compile application
    ${GREEN}📋 list-apps${NC}          List available applications
    
    ${GREEN}❓ help${NC}               Show this help message

${BOLD}COMPILATION PIPELINE:${NC}
    1. P4 → P4HIR MLIR              (p4mlir-translate)    [precompile]
    2. P4HIR → Basic Blocks         (p4hir-partition)     [compile]
    3. Coarse-grained lowering       (p4hir-to-vdrmt/vdpp) [compile]
       - vDRMT: match-action blocks mapped to NIC ASIC pipeline
       - vDPP:  programmable blocks mapped to DPA / ARM
    4a. Fine-grained lowering (DRMT)  (vdrmt-to-bf3drmt)   [compile]
       - DialEgg equality saturation rewrites vDRMT → bf3drmt ops
       - Emits mapper.egg for the deployment optimizer
    4b. Fine-grained lowering (DPA)  (vdpp-to-bf3dpa)     [compile]
       - DialEgg equality saturation rewrites vDPP → bf3dpa ops
       - Appends DPA blocks to mapper.egg
       - Emits dpa_handler.c (FlexIO restricted C) + arm_handler.c
    5. Deployment mapping            (mapper)             [compile]
       - Multi-objective optimizer assigns blocks to hardware
       - Optimizes for latency / throughput / energy
       - Modes: balanced (default), min-latency, max-throughput, min-energy
    6. Control Flow Graph            (graphviz)           [compile]

${BOLD}OPTIONS:${NC}
    MAPPER_MODE=<mode>  Set mapper optimization mode (default: balanced)
                        Values: balanced | min-latency | max-throughput | min-energy
                        Example: MAPPER_MODE=max-throughput $0 compile syn_flood

${BOLD}EXAMPLES:${NC}
    $0 build                                    # Build everything
    $0 precompile-all                           # P4 → MLIR for all apps
    $0 compile syn_flood                        # Full pipeline (balanced)
    MAPPER_MODE=min-latency $0 compile syn_flood  # Optimize for latency
    $0 run syn_flood                            # Build + full compile

${BOLD}WORKFLOW:${NC}
    # One-time setup
    $0 build
    $0 precompile-all

    # Fast iteration
    $0 compile syn_flood          # Uses cached MLIR
    # Edit pass, rebuild
    $0 build-nutcracker
    $0 compile syn_flood          # Instant!

${BOLD}OUTPUT STRUCTURE:${NC}
    apps/<app>/
    ├── <app>.p4              # Source
    └── <app>.mlir            # Precompiled (cached)

    vdsa_output/
    ├── mapper.egg            # Mapper input (all blocks, egglog form)
    ├── mapping.txt           # Mapper output (optimal hw assignment)
    ├── dependency_graph.dot  # Control-flow graph
    ├── dpa_handler.c         # DPA event handler C (FlexIO restricted)
    ├── arm_handler.c         # ARM handler C (standard C99)
    ├── block0/
    │   ├── p4hir.mlir        # Partitioned P4HIR
    │   ├── vdrmt.mlir        # Coarse: vDRMT (match-action)
    │   ├── bf3drmt.mlir      # Fine:   bf3drmt (NIC ASIC ops)
    │   ├── vdpp.mlir         # Coarse: vDPP (programmable)
    │   ├── bf3dpa.mlir       # Fine:   bf3dpa (DPA/ARM ops)
    │   └── metadata.json
    └── block1/ ...

EOF
}

list_apps() {
    print_header "Available Applications"
    
    if [ ! -d "$APPS_DIR" ]; then
        print_error "Apps directory not found: $APPS_DIR"
        exit 1
    fi
    
    for app_dir in "$APPS_DIR"/*; do
        if [ -d "$app_dir" ]; then
            local app_name=$(basename "$app_dir")
            local p4_file="$app_dir/${app_name}.p4"
            
            if [ -f "$p4_file" ]; then
                echo -e "  ${CYAN}📦${NC} ${BOLD}$app_name${NC}"
                
                local desc=$(head -n 5 "$p4_file" | grep -E "^//" | sed 's|^// *||' | head -1)
                [ -n "$desc" ] && echo -e "    $desc"
                
                local mlir_file="$app_dir/${app_name}.mlir"
                if [ -f "$mlir_file" ]; then
                    echo -e "    ${GREEN}✅ Precompiled MLIR available${NC}"
                else
                    echo -e "    ${YELLOW}⚠️  No MLIR (run: precompile $app_name)${NC}"
                fi
                
                echo ""
            fi
        fi
    done
}

compile_all_apps() {
    print_header "Compiling All Applications"
    
    local success_count=0
    local fail_count=0
    
    for app_dir in "$APPS_DIR"/*; do
        if [ -d "$app_dir" ]; then
            local app_name=$(basename "$app_dir")
            local p4_file="$app_dir/${app_name}.p4"
            
            if [ -f "$p4_file" ]; then
                if compile_app "$app_name"; then
                    ((success_count++))
                else
                    ((fail_count++))
                    print_error "Failed to compile $app_name"
                fi
            fi
        fi
    done
    
    print_header "Summary"
    echo -e "  ${GREEN}✅ Success:${NC} $success_count"
    echo -e "  ${RED}❌ Failed:${NC}  $fail_count"
    echo ""
}

# ============================================================================
# Main Entry Point
# ============================================================================

main() {
    cd "$ROOT_DIR"
    
    if [ $# -eq 0 ]; then
        show_usage
        exit 0
    fi
    
    local command=$1
    shift
    
    case "$command" in
        build)
            build_llvm
            build_p4mlir
            build_nutcracker
            setup_p4_includes
            ;;
        
        build-llvm)
            build_llvm
            ;;
        
        build-p4mlir)
            build_p4mlir
            ;;
        
        build-nutcracker)
            build_nutcracker
            ;;
        
        setup)
            setup_p4_includes
            ;;
        
        precompile)
            if [ $# -eq 0 ]; then
                print_error "Application name required"
                echo "Usage: $0 precompile <app_name> [--force]"
                exit 1
            fi
            local force=false
            if [ "$2" = "--force" ]; then
                force=true
            fi
            precompile_app "$1" "$force"
            ;;
        
        precompile-all)
            precompile_all_apps "$@"
            ;;
        
        compile)
            if [ $# -eq 0 ]; then
                print_error "Application name required"
                echo "Usage: $0 compile <app_name> [--target bf2|bf3]"
                exit 1
            fi
            APP_NAME="$1"; shift
            TARGET_BACKEND="bf3"  # default
            while [ $# -gt 0 ]; do
                case "$1" in
                    --target) TARGET_BACKEND="$2"; shift 2 ;;
                    *) shift ;;
                esac
            done
            export TARGET_BACKEND
            compile_app "$APP_NAME"
            ;;
        
        compile-all)
            compile_all_apps
            ;;
        
        run)
            if [ $# -eq 0 ]; then
                print_error "Application name required"
                echo "Usage: $0 run <app_name>"
                exit 1
            fi
            build_nutcracker
            compile_app "$1"
            ;;
        
        clean)
            clean_build
            ;;
        
        clean-mlir)
            clean_mlir
            ;;
        
        clean-all)
            clean_all
            ;;
        
        list-apps)
            list_apps
            ;;
        
        help|--help|-h)
            show_usage
            ;;
        
        *)
            print_error "Unknown command: $command"
            show_usage
            exit 1
            ;;
    esac
}

main "$@"