name: Test and clippy zkVM

on:
  workflow_call:
    inputs:
      zkvm:
        description: 'zkVM to test'
        required: true
        type: string
      cuda:
        description: 'Whether to build CUDA-enabled images'
        required: false
        type: boolean
        default: false
      cuda_archs:
        description: 'Comma-separated CUDA archs to gencode'
        required: false
        type: string
        default: '89,120'
      cluster:
        description: 'Whether to build cluster image'
        required: false
        type: boolean
        default: false
      test_threads:
        description: 'Number of threads when testing via docker'
        required: false
        type: number
      # Remove when we use larger runners, currently only needed to skip for zisk
      skip_prove_test:
        description: 'Whether to skip prove test and ere-dockerized test or not'
        required: false
        type: boolean
        default: false

env:
  CARGO_TERM_COLOR: always

jobs:
  image_meta:
    name: Get image metadata
    runs-on: ubuntu-latest
    outputs:
      dockerfile_changed: ${{ steps.changed_files.outputs.any_changed }}
      image_registry: ${{ steps.image_meta.outputs.image_registry }}
      image_tag: ${{ steps.image_meta.outputs.image_tag }}
      cached_image_tag: ${{ steps.image_meta.outputs.cached_image_tag }}
      base_zkvm_image: ${{ steps.image_meta.outputs.base_zkvm_image }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4
        with:
          ref: ${{ github.event.workflow_run.head_sha || github.sha }}

      - name: Check Dockerfile changes
        id: changed_files
        uses: tj-actions/changed-files@v46
        with:
          files: |
            docker/**
            scripts/**

      - name: Get image metadata
        id: image_meta
        run: |
          if [ "${{ github.event_name }}" == "workflow_run" ]; then
            GIT_SHA="${{ github.event.workflow_run.head_sha }}"
          else
            GIT_SHA="${{ github.sha }}"
          fi
          IMAGE_TAG="${GIT_SHA:0:7}"

          CACHED_IMAGE_TAG=""
          if [ "${{ github.event_name }}" == "pull_request" ] && [ "${{ steps.changed_files.outputs.any_changed }}" == "false" ]; then
            CACHED_IMAGE_TAG="${{ github.event.pull_request.base.sha }}"
            CACHED_IMAGE_TAG="${CACHED_IMAGE_TAG:0:7}"
          fi

          IMAGE_REGISTRY="ghcr.io/${{ github.repository }}"
          BASE_ZKVM_IMAGE="$IMAGE_REGISTRY/ere-base-${{ inputs.zkvm }}:$IMAGE_TAG"

          echo "image_registry=$IMAGE_REGISTRY" >> $GITHUB_OUTPUT
          echo "image_tag=$IMAGE_TAG" >> $GITHUB_OUTPUT
          echo "cached_image_tag=$CACHED_IMAGE_TAG" >> $GITHUB_OUTPUT
          echo "base_zkvm_image=$BASE_ZKVM_IMAGE" >> $GITHUB_OUTPUT

  build_image_cuda_check:
    name: Build image with CUDA enabled
    if: inputs.cuda && github.event_name == 'pull_request'
    needs: image_meta
    runs-on: ubuntu-latest
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4

      - name: Free up disk space
        run: bash .github/scripts/free-up-disk-space.sh

      - name: Pull or build ere-base and ere-base-${{ inputs.zkvm }} images with CUDA enabled
        run: |
          CACHED_TAG="${{ needs.image_meta.outputs.cached_image_tag }}"
          if [ -n "$CACHED_TAG" ]; then
            CACHED_TAG="${CACHED_TAG}-cuda"
          fi
          bash .github/scripts/pull-or-build-base-zkvm-image.sh \
            --zkvm ${{ inputs.zkvm }} \
            --registry ${{ needs.image_meta.outputs.image_registry }} \
            --tag ${{ needs.image_meta.outputs.image_tag }}-cuda \
            --cached-tag "$CACHED_TAG" \
            --cuda-archs '${{ inputs.cuda_archs }}'

      - name: Build ere-server-${{ inputs.zkvm }} image with CUDA enabled
        run: |
          bash .github/scripts/build-image.sh \
            --zkvm ${{ inputs.zkvm }} \
            --registry ${{ needs.image_meta.outputs.image_registry }} \
            --tag ${{ needs.image_meta.outputs.image_tag }}-cuda \
            --server \
            --cuda-archs '${{ inputs.cuda_archs }}'

      - name: Build ere-cluster-${{ inputs.zkvm }} image with CUDA enabled
        if: ${{ inputs.cluster && needs.image_meta.outputs.dockerfile_changed == 'true' }}
        run: |
          bash .github/scripts/build-image.sh \
            --zkvm ${{ inputs.zkvm }} \
            --registry ${{ needs.image_meta.outputs.image_registry }} \
            --tag ${{ needs.image_meta.outputs.image_tag }}-cuda \
            --cluster \
            --cuda-archs '${{ inputs.cuda_archs }}'

  clippy_via_docker:
    name: Clippy via Docker
    needs: image_meta
    runs-on: ubuntu-latest
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4
        with:
          ref: ${{ github.event.workflow_run.head_sha || github.sha }}

      - name: Free up disk space
        run: bash .github/scripts/free-up-disk-space.sh

      - name: Cache dependencies
        uses: Swatinem/rust-cache@v2
        with:
          key: rust-${{ inputs.zkvm }}-${{ hashFiles('Cargo.lock') }}

      - name: Pull published base zkvm image
        if: github.event_name == 'workflow_run'
        run: |
          docker pull ${{ needs.image_meta.outputs.base_zkvm_image }}

      - name: Pull base zkvm image or build locally
        if: github.event_name == 'pull_request'
        run: |
          bash .github/scripts/pull-or-build-base-zkvm-image.sh \
            --zkvm ${{ inputs.zkvm }} \
            --registry ${{ needs.image_meta.outputs.image_registry }} \
            --tag ${{ needs.image_meta.outputs.image_tag }} \
            --cached-tag ${{ needs.image_meta.outputs.cached_image_tag }}

      - name: Run cargo clippy for ere-${{ inputs.zkvm }} via Docker
        run: |
          DOCKER_CMD="docker run \
            --rm \
            --interactive \
            --volume ${{ github.workspace }}:/ere \
            --volume $HOME/.cargo/registry:/usr/local/cargo/registry \
            --volume $HOME/.cargo/git:/usr/local/cargo/git \
            --workdir /ere \
            ${{ needs.image_meta.outputs.base_zkvm_image }} \
            /bin/bash"

          cat <<EOF | $DOCKER_CMD
          set -e -o pipefail

          OPTIONS="--all-targets -- -D warnings"
          cargo clippy --package ere-${{ inputs.zkvm }} \$OPTIONS
          cargo clippy --package ere-compiler --features ${{ inputs.zkvm }} \$OPTIONS
          cargo clippy --package ere-server --features ${{ inputs.zkvm }} \$OPTIONS

          time chown -R $(id -u):$(id -g) /usr/local/cargo/registry
          time chown -R $(id -u):$(id -g) /usr/local/cargo/git
          time chown -R $(id -u):$(id -g) target
          EOF

  test_via_docker:
    name: Test via Docker
    needs: image_meta
    runs-on: ubuntu-latest
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4
        with:
          ref: ${{ github.event.workflow_run.head_sha || github.sha }}

      - name: Free up disk space
        run: bash .github/scripts/free-up-disk-space.sh

      - name: Cache dependencies
        uses: Swatinem/rust-cache@v2
        with:
          key: rust-${{ inputs.zkvm }}-${{ hashFiles('Cargo.lock') }}

      - name: Pull published base zkvm image
        if: github.event_name == 'workflow_run'
        run: |
          docker pull ${{ needs.image_meta.outputs.base_zkvm_image }}

      - name: Pull base zkvm image or build locally
        if: github.event_name == 'pull_request'
        run: |
          bash .github/scripts/pull-or-build-base-zkvm-image.sh \
            --zkvm ${{ inputs.zkvm }} \
            --registry ${{ needs.image_meta.outputs.image_registry }} \
            --tag ${{ needs.image_meta.outputs.image_tag }} \
            --cached-tag ${{ needs.image_meta.outputs.cached_image_tag }}

      - name: Increase swap to 8GB
        run: |
          df -h /

          free -h && swapon --show

          sudo swapoff -a
          sudo fallocate -l 8G /swapfile
          sudo chmod 600 /swapfile
          sudo mkswap /swapfile
          sudo swapon /swapfile

          free -h && swapon --show

      - name: Run cargo test for ere-${{ inputs.zkvm }} via Docker
        run: |
          DOCKER_CMD="docker run \
            --rm \
            --interactive \
            --volume ${{ github.workspace }}:/ere \
            --volume $HOME/.cargo/registry:/usr/local/cargo/registry \
            --volume $HOME/.cargo/git:/usr/local/cargo/git \
            --workdir /ere \
            ${{ needs.image_meta.outputs.base_zkvm_image }} \
            /bin/bash"

          cat <<EOF | $DOCKER_CMD
          set -e -o pipefail

          cargo test --release --package ere-${{ inputs.zkvm }} \
            -- ${{ inputs.skip_prove_test && '--skip prove' || '' }} ${{ inputs.test_threads && format('--test-threads={0}', inputs.test_threads) || '' }}

          time chown -R $(id -u):$(id -g) /usr/local/cargo/registry
          time chown -R $(id -u):$(id -g) /usr/local/cargo/git
          time chown -R $(id -u):$(id -g) target
          EOF

  test_ere_dockerized:
    name: Test ere-dockerized with the selected zkVM
    needs: image_meta
    runs-on: ubuntu-latest
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4
        with:
          ref: ${{ github.event.workflow_run.head_sha || github.sha }}

      - name: Free up disk space
        run: bash .github/scripts/free-up-disk-space.sh

      - name: Install Rust toolchain
        uses: dtolnay/rust-toolchain@master
        with:
          toolchain: 1.88.0

      - name: Cache dependencies
        uses: Swatinem/rust-cache@v2
        with:
          key: rust-${{ inputs.zkvm }}-${{ hashFiles('Cargo.lock') }}

      - name: Pull published images
        if: github.event_name == 'workflow_run'
        run: |
          docker pull ${{ needs.image_meta.outputs.image_registry }}/ere-compiler-${{ inputs.zkvm }}:${{ needs.image_meta.outputs.image_tag }}
          docker pull ${{ needs.image_meta.outputs.image_registry }}/ere-server-${{ inputs.zkvm }}:${{ needs.image_meta.outputs.image_tag }}

      - name: Pull images or build locally
        if: github.event_name == 'pull_request'
        run: |
          bash .github/scripts/pull-or-build-base-zkvm-image.sh \
            --zkvm ${{ inputs.zkvm }} \
            --registry ${{ needs.image_meta.outputs.image_registry }} \
            --tag ${{ needs.image_meta.outputs.image_tag }} \
            --cached-tag ${{ needs.image_meta.outputs.cached_image_tag }}

          # Build ere-compiler-${{ inputs.zkvm }} and ere-server-${{ inputs.zkvm }}
          bash .github/scripts/build-image.sh \
            --zkvm ${{ inputs.zkvm }} \
            --registry ${{ needs.image_meta.outputs.image_registry }} \
            --tag ${{ needs.image_meta.outputs.image_tag }} \
            --compiler \
            --server

          # Reclaim the disk space
          docker image rm ${{ needs.image_meta.outputs.image_registry }}/ere-base:${{ needs.image_meta.outputs.image_tag }} 2>/dev/null || true
          docker image rm ${{ needs.image_meta.outputs.image_registry }}/ere-base-${{ inputs.zkvm }}:${{ needs.image_meta.outputs.image_tag }} 2>/dev/null || true
          docker builder prune -f
          docker image prune -f

      - name: Increase swap to 8GB
        run: |
          df -h /

          free -h && swapon --show

          sudo swapoff -a
          sudo fallocate -l 8G /swapfile
          sudo chmod 600 /swapfile
          sudo mkswap /swapfile
          sudo swapon /swapfile

          free -h && swapon --show

      - name: Run cargo test for ere-${{ inputs.zkvm }} via ere-dockerized
        env:
          RUST_LOG: info
          ERE_IMAGE_REGISTRY: ${{ needs.image_meta.outputs.image_registry }}
        run: |
          cargo test --release --package ere-dockerized \
            -- ${{ inputs.zkvm }} ${{ inputs.skip_prove_test && '--skip prove' || ''  }} --test-threads=1