core/pkgs/build-support/rust/import-cargo-lock.nix

{ fetchgit, fetchurl, lib, writers, python3Packages, runCommand, cargo, jq }:

{
# Cargo lock file
lockFile ? null

  # Cargo lock file contents as string
, lockFileContents ? null

  # Allow `builtins.fetchGit` to be used to not require hashes for git dependencies
, allowBuiltinFetchGit ? false

  # Additional registries to pull sources from
  #   { "https://<registry index URL>" = "https://<registry download URL>"; }
  # where:
  # - "index URL" is the "index" value of the configuration entry for that registry
  #   https://doc.rust-lang.org/cargo/reference/registries.html#using-an-alternate-registry
  # - "download URL" is the "dl" value of its associated index configuration
  #   https://doc.rust-lang.org/cargo/reference/registry-index.html#index-configuration
, extraRegistries ? { }

  # Hashes for git dependencies.
, outputHashes ? { } }@args:

assert (lockFile == null) != (lockFileContents == null);

let
  # Parse a git source into different components.
  parseGit = src:
    let
      parts =
        builtins.match "git\\+([^?]+)(\\?(rev|tag|branch)=(.*))?#(.*)" src;
      type = builtins.elemAt parts 2; # rev, tag or branch
      value = builtins.elemAt parts 3;
    in if parts == null then
      null
    else
      {
        url = builtins.elemAt parts 0;
        sha = builtins.elemAt parts 4;
      } // lib.optionalAttrs (type != null) { inherit type value; };

  # shadows args.lockFileContents
  lockFileContents = if lockFile != null then
    builtins.readFile lockFile
  else
    args.lockFileContents;

  parsedLockFile = builtins.fromTOML lockFileContents;

  packages = parsedLockFile.package;

  # There is no source attribute for the source package itself. But
  # since we do not want to vendor the source package anyway, we can
  # safely skip it.
  depPackages = builtins.filter (p: p ? "source") packages;

  # Create dependent crates from packages.
  #
  # Force evaluation of the git SHA -> hash mapping, so that an error is
  # thrown if there are stale hashes. We cannot rely on gitShaOutputHash
  # being evaluated otherwise, since there could be no git dependencies.
  depCrates =
    builtins.deepSeq gitShaOutputHash (builtins.map mkCrate depPackages);

  # Map package name + version to git commit SHA for packages with a git source.
  namesGitShas = builtins.listToAttrs (builtins.map nameGitSha
    (builtins.filter (pkg: lib.hasPrefix "git+" pkg.source) depPackages));

  nameGitSha = pkg:
    let gitParts = parseGit pkg.source;
    in {
      name = "${pkg.name}-${pkg.version}";
      value = gitParts.sha;
    };

  # Convert the attrset provided through the `outputHashes` argument to a
  # a mapping from git commit SHA -> output hash.
  #
  # There may be multiple different packages with different names
  # originating from the same git repository (typically a Cargo
  # workspace). By using the git commit SHA as a universal identifier,
  # the user does not have to specify the output hash for every package
  # individually.
  gitShaOutputHash = lib.mapAttrs' (nameVer: hash:
    let
      unusedHash = throw
        "A hash was specified for ${nameVer}, but there is no corresponding git dependency.";
      rev = namesGitShas.${nameVer} or unusedHash;
    in {
      name = rev;
      value = hash;
    }) outputHashes;

  # We can't use the existing fetchCrate function, since it uses a
  # recursive hash of the unpacked crate.
  fetchCrate = pkg: downloadUrl:
    let
      checksum =
        pkg.checksum or parsedLockFile.metadata."checksum ${pkg.name} ${pkg.version} (${pkg.source})";
    in assert lib.assertMsg (checksum != null) ''
      Package ${pkg.name} does not have a checksum.
    '';
    fetchurl {
      name = "crate-${pkg.name}-${pkg.version}.tar.gz";
      url = "${downloadUrl}/${pkg.name}/${pkg.version}/download";
      sha256 = checksum;
    };

  registries = {
    "https://github.com/rust-lang/crates.io-index" =
      "https://crates.io/api/v1/crates";
  } // extraRegistries;

  # Replaces values inherited by workspace members.
  replaceWorkspaceValues = writers.writePython3 "replace-workspace-values" {
    libraries = with python3Packages; [ tomli tomli-w ];
    flakeIgnore = [ "E501" "W503" ];
  } (builtins.readFile ./replace-workspace-values.py);

  # Fetch and unpack a crate.
  mkCrate = pkg:
    let
      gitParts = parseGit pkg.source;
      registryIndexUrl = lib.removePrefix "registry+" pkg.source;
    in if lib.hasPrefix "registry+" pkg.source
    && builtins.hasAttr registryIndexUrl registries then
      let crateTarball = fetchCrate pkg registries.${registryIndexUrl};
      in runCommand "${pkg.name}-${pkg.version}" { } ''
        mkdir $out
        tar xf "${crateTarball}" -C $out --strip-components=1

        # Cargo is happy with largely empty metadata.
        printf '{"files":{},"package":"${crateTarball.outputHash}"}' > "$out/.cargo-checksum.json"
      ''
    else if gitParts != null then
      let
        missingHash = throw ''
          No hash was found while vendoring the git dependency ${pkg.name}-${pkg.version}. You can add
          a hash through the `outputHashes` argument of `importCargoLock`:

          outputHashes = {
            "${pkg.name}-${pkg.version}" = "<hash>";
          };

          If you use `buildRustPackage`, you can add this attribute to the `cargoLock`
          attribute set.
        '';
        tree = if gitShaOutputHash ? ${gitParts.sha} then
          fetchgit {
            inherit (gitParts) url;
            rev = gitParts.sha; # The commit SHA is always available.
            sha256 = gitShaOutputHash.${gitParts.sha};
          }
        else if allowBuiltinFetchGit then
          builtins.fetchGit {
            inherit (gitParts) url;
            rev = gitParts.sha;
            allRefs = true;
            submodules = true;
          }
        else
          missingHash;
      in runCommand "${pkg.name}-${pkg.version}" { } ''
        tree=${tree}

        # If the target package is in a workspace, or if it's the top-level
        # crate, we should find the crate path using `cargo metadata`.
        # Some packages do not have a Cargo.toml at the top-level,
        # but only in nested directories.
        # Only check the top-level Cargo.toml, if it actually exists
        if [[ -f $tree/Cargo.toml ]]; then
          crateCargoTOML=$(${cargo}/bin/cargo metadata --format-version 1 --no-deps --manifest-path $tree/Cargo.toml | \
          ${jq}/bin/jq -r '.packages[] | select(.name == "${pkg.name}") | .manifest_path')
        fi

        # If the repository is not a workspace the package might be in a subdirectory.
        if [[ -z $crateCargoTOML ]]; then
          for manifest in $(find $tree -name "Cargo.toml"); do
            echo Looking at $manifest
            crateCargoTOML=$(${cargo}/bin/cargo metadata --format-version 1 --no-deps --manifest-path "$manifest" | ${jq}/bin/jq -r '.packages[] | select(.name == "${pkg.name}") | .manifest_path' || :)
            if [[ ! -z $crateCargoTOML ]]; then
              break
            fi
          done

          if [[ -z $crateCargoTOML ]]; then
            >&2 echo "Cannot find path for crate '${pkg.name}-${pkg.version}' in the tree in: $tree"
            exit 1
          fi
        fi

        echo Found crate ${pkg.name} at $crateCargoTOML
        tree=$(dirname $crateCargoTOML)

        cp -prvL "$tree/" $out
        chmod u+w $out

        if grep -q workspace "$out/Cargo.toml"; then
          chmod u+w "$out/Cargo.toml"
          ${replaceWorkspaceValues} "$out/Cargo.toml" "$(${cargo}/bin/cargo metadata --format-version 1 --no-deps --manifest-path $crateCargoTOML | ${jq}/bin/jq -r .workspace_root)/Cargo.toml"
        fi

        # Cargo is happy with empty metadata.
        printf '{"files":{},"package":null}' > "$out/.cargo-checksum.json"

        # Set up configuration for the vendor directory.
        cat > $out/.cargo-config <<EOF
        [source."${gitParts.url}${
          lib.optionalString (gitParts ? type)
          "?${gitParts.type}=${gitParts.value}"
        }"]
        git = "${gitParts.url}"
        ${lib.optionalString (gitParts ? type)
        ''${gitParts.type} = "${gitParts.value}"''}
        replace-with = "vendored-sources"
        EOF
      ''
    else
      throw "Cannot handle crate source: ${pkg.source}";

  vendorDir = runCommand "cargo-vendor-dir" (if lockFile == null then {
    inherit lockFileContents;
    passAsFile = [ "lockFileContents" ];
  } else {
    passthru = { inherit lockFile; };
  }) ''
        mkdir -p $out/.cargo

        ${
          if lockFile != null then
            "ln -s ${lockFile} $out/Cargo.lock"
          else
            "cp $lockFileContentsPath $out/Cargo.lock"
        }

        cat > $out/.cargo/config <<EOF
    [source.crates-io]
    replace-with = "vendored-sources"

    [source.vendored-sources]
    directory = "cargo-vendor-dir"
    EOF

        declare -A keysSeen

        for registry in ${toString (builtins.attrNames extraRegistries)}; do
          cat >> $out/.cargo/config <<EOF

    [source."$registry"]
    registry = "$registry"
    replace-with = "vendored-sources"
    EOF
        done

        for crate in ${toString depCrates}; do
          # Link the crate directory, removing the output path hash from the destination.
          ln -s "$crate" $out/$(basename "$crate" | cut -c 34-)

          if [ -e "$crate/.cargo-config" ]; then
            key=$(sed 's/\[source\."\(.*\)"\]/\1/; t; d' < "$crate/.cargo-config")
            if [[ -z ''${keysSeen[$key]} ]]; then
              keysSeen[$key]=1
              cat "$crate/.cargo-config" >> $out/.cargo/config
            fi
          fi
        done
  '';
in vendorDir