nix-servo/flake.nix

{
  description = "Firmware for Sinara Fast-Servo based on Not-OS and Linien";
  
  inputs.nixpkgs.url = github:NixOS/nixpkgs/nixos-24.11;
  inputs.not-os.url = github:cleverca22/not-os;
  inputs.not-os.inputs.nixpkgs.follows = "nixpkgs";

  inputs.src-migen = { url = github:m-labs/migen; flake = false; };
  inputs.src-misoc = { type = "git"; url = "https://github.com/m-labs/misoc.git"; submodules = true; flake = false; };

  outputs = { self, nixpkgs, not-os, src-migen, src-misoc }:
    let
      pkgs = import nixpkgs { system = "x86_64-linux"; overlays = [ crosspkgs-overlay ]; };
      pkgs-armv7l = pkgs.pkgsCross.zynq-armv7l-linux;
      fsbl-support = ./fast-servo/fsbl-support;

      version = "2.1.0";
      linien-src = pkgs.fetchFromGitHub {
          owner = "linien-org";
          repo = "linien";
          rev = "v" + version;
          sha256 = "sha256-j6oiP/usLfV5HZtKLcXQ5pHhhxRG05kP2FMwingiWm0=";
      };

      patched-not-os = pkgs.applyPatches {
        name = "not-os-patched";
        src = not-os;
        patches = [
          ./not-os-patches/network.patch
          ./not-os-patches/pr-28.patch
          ./not-os-patches/pr-29.patch
          ./not-os-patches/pr-30.patch
          ./not-os-patches/pr-31.patch
          ./not-os-patches/pr-33.patch
          ./not-os-patches/iproute2.patch
        ];
      };

      crossSystem = {
        system = "armv7l-linux";
        linux-kernel = {
          name = "zynq";
          baseConfig = "multi_v7_defconfig";
          target = "uImage";
          installTarget = "uImage";
          autoModules = false;
          DTB = true;
          makeFlags = [ "LOADADDR=0x8000" ];
        };
      };

      crosspkgs-overlay = (self: super: {
        pkgsCross = super.pkgsCross // {
          zynq-baremetal = import super.path {
            system = "x86_64-linux";
            crossSystem = {
              config = "arm-none-eabihf";
              libc = "newlib";
              gcc.cpu = "cortex-a9";
              gcc.fpu = "vfpv3";
            };
          };
          zynq-armv7l-linux = import super.path {
            system = "x86_64-linux";
            inherit crossSystem;
          };
        };
      });

      migen = pkgs.python3Packages.buildPythonPackage rec {
        name = "migen";
        src = src-migen;
        format = "pyproject";
        nativeBuildInputs = [ pkgs.python3Packages.setuptools ];
        propagatedBuildInputs = [ pkgs.python3Packages.colorama ];
      };

      misoc = pkgs.python3Packages.buildPythonPackage {
        name = "misoc";
        src = src-misoc;
        propagatedBuildInputs = with pkgs.python3Packages; [ jinja2 numpy migen pyserial asyncserial ];
      };

      vivado = pkgs.buildFHSEnv {
        name = "vivado";
        targetPkgs = pkgs: with pkgs; let
        # Apply patch from https://github.com/nix-community/nix-environments/pull/54
        # to fix ncurses libtinfo.so's soname issue
        ncurses' = ncurses5.overrideAttrs (old: {
          configureFlags = old.configureFlags ++ [ "--with-termlib" ];
          postFixup = "";
        });
        in [
          libxcrypt-legacy
          (ncurses'.override { unicodeSupport = false; })
          zlib
          libuuid
          xorg.libSM
          xorg.libICE
          xorg.libXrender
          xorg.libX11
          xorg.libXext
          xorg.libXtst
          xorg.libXi
          freetype
          fontconfig
        ];
        profile = "set -e; source /opt/Xilinx/Vivado/2024.2/settings64.sh";
        runScript = "vivado";
      };

      cma = pkgs-armv7l.python3Packages.buildPythonPackage rec {
        pname = "cma";
        version = "3.3.0";
        src = pkgs.fetchFromGitHub {
          owner = "CMA-ES";
          repo = "pycma";
          rev = "refs/tags/r${version}";
          hash = "sha256-+UJI3hDVbDMfRF4bkwHED3eJCHzxS2hO4YPUzJqcoQI=";
        };

        propagatedBuildInputs = [ pkgs-armv7l.python3Packages.numpy ];

        pythonImportsCheck = [ "cma" ];

        checkPhase = ''
          # At least one doctest fails, thus only limited amount of files is tested
          python -m cma.test interfaces.py purecma.py logger.py optimization_tools.py transformations.py
        '';
      };

      pyrp3 = pkgs-armv7l.python3Packages.buildPythonPackage rec {
        pname = "pyrp3";
        version = "2.1.0";
        format = "pyproject";
        src = pkgs.fetchFromGitHub {
          owner = "linien-org";
          repo = "pyrp3";
          rev = "v${version}";
          hash = "sha256-ol1QGXyCOei94iIPIocuTRHBxa5jKSH5RzjzROfZaBI=";
        };
        patches = ./fast-servo/linien-pyrp3-monitor.patch;
        nativeBuildInputs = [
          pkgs-armv7l.python3Packages.setuptools
          pkgs-armv7l.gcc
        ];
        postInstall = ''
          cp monitor/libmonitor.so $out/lib
        '';
        postFixup = ''
          substituteInPlace $out/${pkgs.python3.sitePackages}/pyrp3/raw_memory.py \
            --replace "libmonitor.so" "$out/lib/libmonitor.so"
        '';
        propagatedBuildInputs = with pkgs-armv7l.python3Packages; [
          cached-property
          numpy
          rpyc
        ];
      };

      linien-client = pkgs.python3Packages.buildPythonPackage rec {
        pname = "linien-client";
        inherit version;
        src = linien-src;
        
        pyproject = true;

        sourceRoot = "${src.name}/linien-client";

        preBuild = ''
          export HOME=$(mktemp -d)
        '';

        nativeBuildInputs = [ pkgs.python3Packages.setuptools ];

        patches = [
          ./fast-servo/linien-client-ssh-port-change.patch
        ];

        doInstallCheck = false;
        doCheck = false;
        propagatedBuildInputs = with pkgs.python3Packages; [
          fabric
          typing-extensions
          linien-common
        ];

        pythonImportsCheck = [ "linien_client" ];

        };

      linien-gui = pkgs.python3Packages.buildPythonApplication rec {
        pname = "linien-gui";
        inherit version;
        src = linien-src;
        pyproject = true;

        sourceRoot = "${src.name}/linien-gui";
        nativeBuildInputs = with pkgs.python3Packages; [
          setuptools
        ] ++ [
          pkgs.qt5.wrapQtAppsHook
        ];

        # Makes qt-wayland appear in the qt paths injected by the wrapper - helps users
        # with `QT_QPA_PLATFORM=wayland` in their environment.
        buildInputs = [
          pkgs.qt5.qtwayland
        ];

        propagatedBuildInputs = with pkgs.python3Packages; [
          click
          pyqtgraph
          pyqt5
          requests
          superqt
        ] ++ [ linien-client ];

        dontWrapQtApps = true;

        preFixup = ''
          makeWrapperArgs+=("''${qtWrapperArgs[@]}")
        '';
      };

      linien-server = pkgs-armv7l.python3Packages.buildPythonPackage rec {
        pname = "linien-server";
        inherit version;
        src = linien-src;
        pyproject = true;

        sourceRoot = "${src.name}/linien-server";
        patches = [
          ./fast-servo/linien-server-fast-servo.patch
        ];

        postPatch = ''
          cp ${fast-servo-gateware}/csrmap.py linien_server/csrmap.py
          substituteInPlace linien_server/acquisition.py \
            --replace "  start_nginx()" "" \
            --replace "  stop_nginx()" "" \
            --replace "  flash_fpga()" ""
        '';
        nativeBuildInputs = [ pkgs-armv7l.python3Packages.setuptools ];
        propagatedBuildInputs = with pkgs-armv7l.python3Packages; [
          fire
          influxdb-client
          pylpsd
          linien-common
        ] ++ [ 
          cma
          pyrp3
        ];
      };

      fast-servo-gateware = pkgs.stdenv.mkDerivation rec {
        name = "fast-servo-gateware";
        inherit (pkgs.python3Packages.linien-common) src;
        prePatch = ''
          mkdir -p fast_servo/gateware
          cp -r ${./fast-servo/linien-gateware}/. fast_servo/gateware
        '';
        patches = [
          fast-servo/linien-gateware-fast-servo.patch 
          fast-servo/autolock_pipeline.patch
          fast-servo/iir_pipeline.patch
          fast-servo/linien_module_pipeline.patch
          fast-servo/pid_pipeline.patch
          fast-servo/pid_err_sig_pipeline.patch
        ];
        nativeBuildInputs = [
          (pkgs.python3.withPackages(ps: [
            migen misoc
            (ps.linien-common.overrideAttrs(oa: {
              # config.py tries to access $HOME, but we do not need it for building gateware
              postPatch = ''
                echo > linien_common/config.py
                echo > linien_common/__init__.py
              '';
              doCheck = false;
              }))
            ]))
          vivado
        ];
        buildPhase = ''
          python -m gateware.fpga_image_helper -p fastservo
        '';
        installPhase = ''
          mkdir -p $out $out/nix-support
          cp gateware/build/top.bit $out
          cp linien-server/linien_server/gateware.bin $out
          cp linien-server/linien_server/csrmap.py $out
          echo file binary-dist $out/top.bit >> $out/nix-support/hydra-build-products
          echo file binary-dist $out/gateware.bin >> $out/nix-support/hydra-build-products
        '';
      };

      pyfastservo = pkgs-armv7l.python3Packages.buildPythonPackage rec {
        name = "pyfastservo";
        src = ./fast-servo;
        preBuild = ''
          cat > setup.py << EOF
          from setuptools import setup

          setup(
            name="pyfastservo",
            packages=["pyfastservo"],
            install_requires=["spidev", "smbus2"],
            entry_points = {"console_scripts": ["fp_leds=pyfastservo.fp_leds:main"]},
          )
          EOF
        '';
        propagatedBuildInputs = with pkgs-armv7l.python3Packages; [
          spidev
          smbus2
        ];
      };

      mkbootimage = pkgs.stdenv.mkDerivation {
        pname = "mkbootimage";
        version = "2.3dev";

        src = pkgs.fetchFromGitHub {
          owner = "antmicro";
          repo = "zynq-mkbootimage";
          rev = "872363ce32c249f8278cf107bc6d3bdeb38d849f";
          sha256 = "sha256-5FPyAhUWZDwHbqmp9J2ZXTmjaXPz+dzrJMolaNwADHs=";
        };

        propagatedBuildInputs = [ pkgs.libelf pkgs.pcre ];
        patchPhase = ''
          substituteInPlace Makefile --replace "git rev-parse --short HEAD" "echo nix"
        '';
        installPhase = ''
          mkdir -p $out/bin
          cp mkbootimage $out/bin
        '';
        # fix crash; see https://github.com/xmrig/xmrig/issues/3305
        hardeningDisable = [ "fortify" ];
      };

      board-package-set = { board }: let
        not-os-configured = (import patched-not-os {
          inherit nixpkgs;
          extraModules = [
            "${patched-not-os}/zynq_image.nix"
          ] ++ pkgs.lib.optionals (board == "fast-servo") [
            ({ config, pkgs, lib, ... }: {
              environment.systemPackages = [
                linien-server
                (pkgs.python3.withPackages(ps: [ pyfastservo ]))
              ];
              boot.postBootCommands = lib.mkAfter ''

                # Program the FPGA
                set +x
                echo "Loading bitstream into SRAM..."
                echo 0 > /sys/class/fpga_manager/fpga0/flags
                mkdir -p /lib/firmware
                cp ${fast-servo-gateware}/gateware.bin /lib/firmware/
                echo gateware.bin > /sys/class/fpga_manager/fpga0/firmware

                # Run device init scripts
                echo "Initializing clock generator, ADC, and DAC..."
                python3 -m pyfastservo.initialize
              '';
            })];
          system = "x86_64-linux";
          inherit crossSystem;
        });

        not-os-build = not-os-configured.config.system.build;

        fsbl = pkgs.stdenv.mkDerivation {
          name = "${board}-fsbl";
          src = pkgs.fetchFromGitHub {
            owner = "Xilinx";
            repo = "embeddedsw";
            rev = "xilinx_v2022.2";
            sha256 = "sha256-UDz9KK/Hw3qM1BAeKif30rE8Bi6C2uvuZlvyvtJCMfw=";
          };
          nativeBuildInputs = [
            pkgs.pkgsCross.zynq-baremetal.buildPackages.binutils
            pkgs.pkgsCross.zynq-baremetal.buildPackages.gcc
          ];
          postUnpack = ''
            mkdir -p $sourceRoot/lib/sw_apps/zynq_fsbl/misc/fast-servo
            cp $sourceRoot/lib/sw_apps/zynq_fsbl/misc/zc706/* $sourceRoot/lib/sw_apps/zynq_fsbl/misc/fast-servo
            cp ${fsbl-support}/* $sourceRoot/lib/sw_apps/zynq_fsbl/misc/fast-servo
          '';
          patches = [] ++ pkgs.lib.optional (board == "fast-servo") ./fast-servo/fsbl.patch;
          postPatch = ''
            patchShebangs lib/sw_apps/zynq_fsbl/misc/copy_bsp.sh

            for x in lib/sw_apps/zynq_fsbl/src/Makefile lib/sw_apps/zynq_fsbl/misc/copy_bsp.sh lib/bsp/standalone/src/arm/cortexa9/gcc/Makefile; do
              substituteInPlace $x \
                --replace "arm-none-eabi-" "arm-none-eabihf-"
            done
          '';
          buildPhase = ''
            cd lib/sw_apps/zynq_fsbl/src
            make BOARD=${board} "CFLAGS=-DFSBL_DEBUG_INFO -g"
          '';
          installPhase = ''
            mkdir $out
            cp fsbl.elf $out
          '';
          doCheck = false;
          dontFixup = true;
        };

        u-boot = (pkgs-armv7l.buildUBoot {
            name = "${board}-u-boot";
            defconfig = "xilinx_zynq_virt_defconfig";
            patches = [] ++ pkgs.lib.optional (board == "fast-servo") ./fast-servo/u-boot.patch;
            preConfigure = ''
              export DEVICE_TREE=zynq-${board}
            '';
            extraConfig = ''
              CONFIG_SYS_PROMPT="${board}-boot> "
              CONFIG_AUTOBOOT=y
              CONFIG_BOOTCOMMAND="${builtins.replaceStrings [ "\n" ] [ "; " ] ''
                setenv bootargs 'root=/dev/mmcblk0p2 console=ttyPS0,115200n8 systemConfig=${builtins.unsafeDiscardStringContext not-os-build.toplevel}'
                fatload mmc 0 0x6400000 uImage
                fatload mmc 0 0x8000000 ${board}.dtb
                fatload mmc 0 0xA400000 uRamdisk.image.gz
                bootm 0x6400000 0xA400000 0x8000000
              ''}"
              CONFIG_BOOTDELAY=0
              CONFIG_USE_BOOTCOMMAND=y
            '';
            extraMeta.platforms = [ "armv7l-linux" ];
            filesToInstall = [ "u-boot.elf" ];
          }).overrideAttrs (oldAttrs: {
            postUnpack = ''
              cp ${fast-servo/fast-servo.dts} $sourceRoot/arch/arm/dts/zynq-fast-servo.dts
            '';
            postInstall = ''
              mkdir -p $out/dts
              cp arch/arm/dts/zynq-fast-servo.dts $out/dts
              cp arch/arm/dts/zynq-zc706.dts $out/dts
              cp arch/arm/dts/zynq-7000.dtsi $out/dts
            '';
          });

        bootimage = pkgs.runCommand "${board}-bootimage"
          {
            buildInputs = [ mkbootimage ];
          }
          ''
            bifdir=`mktemp -d`
            cd $bifdir
            ln -s ${fsbl}/fsbl.elf fsbl.elf
            ln -s ${u-boot}/u-boot.elf u-boot.elf
            cat > boot.bif << EOF
            the_ROM_image:
            {
              [bootloader]fsbl.elf
              u-boot.elf
            }
            EOF
            mkdir $out $out/nix-support
            mkbootimage boot.bif $out/boot.bin
            echo file binary-dist $out/boot.bin >> $out/nix-support/hydra-build-products
          '';

        dtb = pkgs.runCommand "${board}-dtb"
          {
            buildInputs = [ pkgs.gcc pkgs.dtc ];
          }
          ''
            mkdir -p $out
            cp ${u-boot}/dts/zynq-${board}.dts .

            if [ ${board} == "zc706" ]; then
              mv zynq-${board}.dts zynq-${board}-top.dts
              cp ${u-boot}/dts/zynq-7000.dtsi .
              gcc -E -nostdinc -undef -D__DTS__ -x assembler-with-cpp -o zynq-${board}.dts zynq-${board}-top.dts
            fi

            dtc -I dts -O dtb -o ${board}.dtb zynq-${board}.dts
            cp ${board}.dtb $out
          '';

        sd-image = let
          rootfsImage = pkgs.callPackage (pkgs.path + "/nixos/lib/make-ext4-fs.nix") {
            storePaths = [ not-os-build.toplevel ];
            volumeLabel = "ROOT";
          };
          # Current firmware (kernel, bootimage, etc..) takes ~18MB
          firmwareSize = 30;
          firmwarePartitionOffset = 8;
          in pkgs.stdenv.mkDerivation {
            name = "${board}-sd-image";
            nativeBuildInputs = with pkgs; [ dosfstools mtools libfaketime util-linux parted ];
            buildCommand = ''
              mkdir -p $out/nix-support $out/sd-image
              export img=$out/sd-image/sd-image.img

              echo "${pkgs.stdenv.buildPlatform.system}" > $out/nix-support/system
              echo "file sd-image $img" >> $out/nix-support/hydra-build-products

              gap=${toString firmwarePartitionOffset}

              rootSizeBlocks=$(du -B 512 --apparent-size ${rootfsImage} | awk '{ print $1 }')
              firmwareSizeBlocks=$((${toString firmwareSize} * 1024 * 1024 / 512))
              imageSize=$((rootSizeBlocks * 512 + firmwareSizeBlocks * 512 + gap * 1024 * 1024))
              truncate -s $imageSize $img

              fat32Start="$((gap))MB"
              fat32End="$((gap + ${toString firmwareSize}))MB"

              parted $img mklabel msdos
              parted $img mkpart primary fat32 $fat32Start $fat32End
              parted $img mkpart primary ext4 $fat32End 100%
              parted $img set 1 boot on

              eval $(partx $img -o START,SECTORS --nr 2 --pairs)
              dd conv=notrunc if=${rootfsImage} of=$img seek=$START count=$SECTORS

              eval $(partx $img -o START,SECTORS --nr 1 --pairs)
              truncate -s $((SECTORS * 512)) firmware_part.img
              faketime "1970-01-01 00:00:00" mkfs.vfat -n BOOT firmware_part.img

              mkdir firmware
              cp ${bootimage}/boot.bin firmware/
              cp ${dtb}/${board}.dtb firmware/
              cp ${not-os-build.kernel}/uImage firmware/
              cp ${not-os-build.uRamdisk}/initrd firmware/uRamdisk.image.gz

              (cd firmware; mcopy -psvm -i ../firmware_part.img ./* ::)
              dd conv=notrunc if=firmware_part.img of=$img seek=$START count=$SECTORS
            '';
          };

        not-os-qemu = let
          qemuScript = ''
            #!/bin/bash
            export PATH=${pkgs.qemu}/bin:$PATH
            IMGDIR=$(mktemp -d /tmp/not-os-qemu-XXXXXX)
            BASE=$(realpath $(dirname $0))

            qemu-img convert -O qcow2 -f raw -o preallocation=metadata $BASE/sd-image.img $IMGDIR/sd-sparse.qcow2
            qemu-img create -F qcow2 -f qcow2 -b $IMGDIR/sd-sparse.qcow2 $IMGDIR/sd-overlay.qcow2 2G

            # Some command arguments are based from samples in Xilinx QEMU User Documentation
            # See: https://xilinx-wiki.atlassian.net/wiki/spaces/A/pages/821854273/Running+Bare+Metal+Applications+on+QEMU

            qemu-system-arm \
              -M xilinx-zynq-a9 \
              -m 1024 \
              $([ ${board} = "zc706" ] && echo "-serial /dev/null") -serial stdio \
              -display none \
              -kernel $BASE/u-boot.elf \
              -sd $IMGDIR/sd-overlay.qcow2

            rm -rf $IMGDIR
          '';
          in pkgs.runCommand "${board}-qemu" {
            inherit qemuScript;
            passAsFile = [ "qemuScript" ];
            preferLocalBuild = true;
          }
          ''
            mkdir $out
            cd $out
            cp -s ${u-boot}/u-boot.elf .
            cp -s ${sd-image}/sd-image/sd-image.img .
            cp $qemuScriptPath qemu-script
            chmod +x qemu-script
            patchShebangs qemu-script
          '';
      in {
        "${board}-fsbl" = fsbl;
        "${board}-u-boot" = u-boot;
        "${board}-bootimage" = bootimage;
        "${board}-dtb" = dtb;
        "${board}-sd-image" = sd-image;
        "${board}-qemu" = not-os-qemu;
      };
    in rec {
      devShell.x86_64-linux = pkgs.mkShell {
        name = "nix-servo-dev_shell";
        buildInputs = with pkgs.python3Packages; [
          linien-common 
          matplotlib
        ] ++ [ linien-client linien-gui ];
      };

      packages.x86_64-linux = {
        inherit mkbootimage;
        inherit migen misoc vivado;
      };
      packages.armv7l-linux = {
        inherit fast-servo-gateware linien-server;
      } //
      (board-package-set { board = "zc706"; }) //
      (board-package-set { board = "fast-servo"; });
      hydraJobs = packages.x86_64-linux // packages.armv7l-linux;
    };

  nixConfig = {
    allow-import-from-derivation = true;
  };
}