flake.nix

{
  description = "rules_mojo";

  nixConfig = {
    bash-prompt-prefix = "(rules_mojo) ";
    bash-prompt = ''\[\033]0;\u@\h:\w\007\]\[\033[01;32m\]\u@\h\[\033[01;34m\] \w \$\[\033[00m\]'';
    bash-prompt-suffix = " ";
  };

  inputs = {
    nixpkgs = {
      url = "github:nixos/nixpkgs";

      # This needs to follow the `nixpkgs` from nativelink so that the local LRE
      # toolchains are in sync with the remote toolchains.
      follows = "nativelink/nixpkgs";
    };
    flake-utils.url = "github:numtide/flake-utils";
    flake-parts = {
      url = "github:hercules-ci/flake-parts";
      inputs.nixpkgs-lib.follows = "nixpkgs";
    };
    pre-commit-hooks = {
      url = "github:cachix/pre-commit-hooks.nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
    nativelink = {
      # Note: Keep this commit in sync with the LRE commit in `MODULE.bazel`.
      url = "github:TraceMachina/nativelink/2adda2475eed578d610a66b98f965922656061af";

      # This repository provides the autogenerated LRE toolchains which are
      # dependent on the nixpkgs version in the nativelink repository. To keep
      # the local LRE toolchains aligned with remote LRE, we need to use the
      # nixpkgs used by nativelink as the the "global" nixpkgs. We do this by
      # setting `nixpkgs.follows = "nativelink/nixpkgs"` above.

      inputs = {
        flake-utils.follows = "flake-utils";
        flake-parts.follows = "flake-parts";
        git-hooks.follows = "pre-commit-hooks";
      };
    };
    rules_ll = {
      url = "github:eomii/rules_ll/a8a26cc78c43896b179caaa263857bd6c88a1da9";
      inputs = {
        nixpkgs.follows = "nixpkgs";
        flake-utils.follows = "flake-utils";
        flake-parts.follows = "flake-parts";
        git-hooks.follows = "pre-commit-hooks";
        nativelink.follows = "nativelink";
      };
    };
    nix2container = {
      follows = "nativelink/nix2container";
    };
  };

  outputs =
    { self
    , nixpkgs
    , pre-commit-hooks
    , flake-parts
    , nativelink
    , rules_ll
    , nix2container
    , ...
    } @ inputs:
    flake-parts.lib.mkFlake { inherit inputs; }
      {
        systems = [
          "x86_64-linux"
        ];
        imports = [
          inputs.nativelink.flakeModule
          inputs.pre-commit-hooks.flakeModule
          inputs.rules_ll.flakeModule
          ./flake-module.nix
        ];
        perSystem =
          { config
          , pkgs
          , system
          , lib
          , ...
          }:
          let
            inherit (inputs.nativelink.packages.${system}) native-cli lre-cc nativelink-debug;

            tinfo6-bin = pkgs.callPackage ./tinfo6-bin.nix { };

            ncurses6-bin = pkgs.callPackage ./ncurses6-bin.nix { };

            lre-mojo-cluster = pkgs.callPackage ./local-remote-execution/lre-mojo-cluster.nix {
              inherit native-cli;
            };

            inherit (nix2container.packages.${system}.nix2container) buildImage;

            mojo = pkgs.callPackage ./mojo.nix {
              inherit (pkgs.lib) makeBinPath makeLibraryPath;
              inherit (pkgs.llvmPackages_18) clang;
              ncurses = ncurses6-bin;
              tinfo = tinfo6-bin;
            };

            lre-mojo = pkgs.callPackage ./local-remote-execution/lre-mojo.nix {
              mojoEnv = self.lib.defaultMojoEnv {
                inherit pkgs mojo;
              };
              inherit buildImage lre-cc;
            };

            lre-kill-the-mojo = pkgs.callPackage ./local-remote-execution/lre-kill-the-mojo.nix { };

            createWorker = pkgs.callPackage ./local-remote-execution/create-worker.nix {
              inherit buildImage self;
              nativelink = nativelink-debug;
            };

            docs = pkgs.writeShellScriptBin "docs" ''
              SRC_ROOT=$(git rev-parse --show-toplevel)

              bazel build @rules_mojo//mojo:docs || exit
              chmod 644 $SRC_ROOT/bazel-bin/mojo/*.md
              cp $SRC_ROOT/bazel-bin/mojo/*.md $SRC_ROOT/docs/src/content/docs/reference/
              cp $SRC_ROOT/bazel-bin/mojo/defs.md $SRC_ROOT/docs/src/content/docs/rules/

              # Rerun the pre-commit hooks so that we do not need to stage everything twice.
              git add docs
              pre-commit run --all-files
              git add docs

              # Technically unnecessary, but we want to show users whether all tests pass.
              printf '
              **************************
              RERUNNING PRE-COMMIT HOOKS
              **************************
              '
              pre-commit run --all-files
            '';

            bazel = pkgs.writeShellScriptBin "bazel" ''
              unset TMPDIR TMP
              exec ${pkgs.bazelisk}/bin/bazelisk "$@"
            '';

            lre-bazel = pkgs.callPackage ./lre-bazel.nix { inherit bazel; };
          in
          {
            _module.args.pkgs = import self.inputs.nixpkgs {
              inherit system;
              # CUDA support
              # config.allowUnfree = true;
              # config.cudaSupport = true;
            };
            local-remote-execution.settings = {
              inherit (lre-cc.meta) Env;
            };
            pre-commit.settings = {
              hooks = import ./pre-commit-hooks.nix { inherit pkgs; };
            };
            rules_ll.settings.llEnv =
              let
                openssl = pkgs.openssl.override { static = true; };
              in
              rules_ll.lib.defaultLlEnv {
                inherit pkgs;
                LL_CFLAGS = "-I${openssl.dev}/include";
                LL_LDFLAGS = "-L${openssl.out}/lib";
              };
            rules_mojo.settings.mojoEnv = self.lib.defaultMojoEnv {
              inherit pkgs mojo;
            };
            packages = {
              inherit lre-mojo lre-cc mojo lre-mojo-cluster lre-kill-the-mojo bazel lre-bazel;
              nativelink-worker-lre-mojo = createWorker lre-mojo;
            };
            devShells.default = pkgs.mkShell {
              nativeBuildInputs = [
                bazel
                docs
                lre-bazel
                lre-kill-the-mojo
                lre-mojo-cluster
                mojo
                pkgs.kubectl
                pkgs.python312
                pkgs.tektoncd-cli
                pkgs.kind
                pkgs.vale
                pkgs.bun
              ];

              shellHook = ''
                # Generate the .pre-commit-config.yaml symlink when entering the
                # development shell.
                ${config.pre-commit.installationScript}

                # Generate .bazelrc.ll which containes action-env
                # configuration when rules_ll is run from a nix environment.
                ${config.rules_ll.installationScript}

                # Generate .bazelrc.mojo which contains Bazel configuration
                # when rules_mojo is run from a nix environment.
                ${config.rules_mojo.installationScript}

                # Generate .bazelrc.lre which configures the LRE toolchains.
                ${config.local-remote-execution.installationScript}

                # Ensure that the bazel command points to our custom wrapper.
                [[ $(type -t bazel) == "alias" ]] && unalias bazel
              '';
            };
          };
      } // {
      flakeModule = ./flake-module.nix;
      lib = { defaultMojoEnv = import ./modules/defaultMojoEnv.nix; };
    };
}