nixpkgs/pkgs/development/cuda-modules/generic-builders/multiplex.nix

{
  lib,
  cudaLib,
  cudaMajorMinorVersion,
  redistSystem,
  stdenv,
  # Builder-specific arguments
  # Short package name (e.g., "cuda_cccl")
  # pname : String
  pname,
  # Common name (e.g., "cutensor" or "cudnn") -- used in the URL.
  # Also known as the Redistributable Name.
  # redistName : String,
  redistName,
  # releasesModule :: Path
  # A path to a module which provides a `releases` attribute
  releasesModule,
  # shims :: Path
  # A path to a module which provides a `shims` attribute
  # The redistribRelease is only used in ./manifest.nix for the package version
  # and the package description (which NVIDIA's manifest calls the "name").
  # It's also used for fetching the source, but we override that since we can't
  # re-use that portion of the functionality (different URLs, etc.).
  # The featureRelease is used to populate meta.platforms (by way of looking at the attribute names), determine the
  # outputs of the package, and provide additional package-specific constraints (e.g., min/max supported CUDA versions,
  # required versions of other packages, etc.).
  # shimFn :: {package, redistSystem} -> AttrSet
  shimsFn ? (throw "shimsFn must be provided"),
}:
let
  evaluatedModules = lib.modules.evalModules {
    modules = [
      ../modules
      releasesModule
    ];
  };

  # NOTE: Important types:
  # - Releases: ../modules/${pname}/releases/releases.nix
  # - Package: ../modules/${pname}/releases/package.nix

  # Check whether a package supports our CUDA version.
  # satisfiesCudaVersion :: Package -> Bool
  satisfiesCudaVersion =
    package:
    lib.versionAtLeast cudaMajorMinorVersion package.minCudaVersion
    && lib.versionAtLeast package.maxCudaVersion cudaMajorMinorVersion;

  # FIXME: do this at the module system level
  propagatePlatforms = lib.mapAttrs (redistSystem: lib.map (p: { inherit redistSystem; } // p));

  # Releases for all platforms and all CUDA versions.
  allReleases = propagatePlatforms evaluatedModules.config.${pname}.releases;

  # Releases for all platforms and our CUDA version.
  allReleases' = lib.mapAttrs (_: lib.filter satisfiesCudaVersion) allReleases;

  # Packages for all platforms and our CUDA versions.
  allPackages = lib.concatLists (lib.attrValues allReleases');

  packageOlder = p1: p2: lib.versionOlder p1.version p2.version;
  packageSupportedPlatform = p: p.redistSystem == redistSystem;

  # Compute versioned attribute name to be used in this package set
  # Patch version changes should not break the build, so we only use major and minor
  # computeName :: Package -> String
  computeName = { version, ... }: cudaLib.mkVersionedName pname (lib.versions.majorMinor version);

  # The newest package for each major-minor version, with newest first.
  # newestPackages :: List Package
  newestPackages =
    let
      newestForEachMajorMinorVersion = lib.foldl' (
        newestPackages: package:
        let
          majorMinorVersion = lib.versions.majorMinor package.version;
          existingPackage = newestPackages.${majorMinorVersion} or null;
        in
        newestPackages
        // {
          ${majorMinorVersion} =
            # Only keep the existing package if it is newer than the one we are considering or it is supported on the
            # current platform and the one we are considering is not.
            if
              existingPackage != null
              && (
                packageOlder package existingPackage
                || (!packageSupportedPlatform package && packageSupportedPlatform existingPackage)
              )
            then
              existingPackage
            else
              package;
        }
      ) { } allPackages;
    in
    # Sort the packages by version so the newest is first.
    # NOTE: builtins.sort requires a strict weak ordering, so we must use versionOlder rather than versionAtLeast.
    # See https://github.com/NixOS/nixpkgs/commit/9fd753ea84e5035b357a275324e7fd7ccfb1fc77.
    lib.sort (lib.flip packageOlder) (lib.attrValues newestForEachMajorMinorVersion);

  extension =
    final: _:
    let
      # Builds our package into derivation and wraps it in a nameValuePair, where the name is the versioned name
      # of the package.
      buildPackage =
        package:
        let
          shims = final.callPackage shimsFn { inherit package redistSystem; };
          name = computeName package;
          drv = final.callPackage ./manifest.nix {
            inherit pname redistName;
            inherit (shims) redistribRelease featureRelease;
          };
        in
        lib.nameValuePair name drv;

      # versionedDerivations :: AttrSet Derivation
      versionedDerivations = builtins.listToAttrs (lib.map buildPackage newestPackages);

      defaultDerivation = {
        ${pname} = (buildPackage (lib.head newestPackages)).value;
      };
    in
    # NOTE: Must condition on the length of newestPackages to avoid non-total function lib.head aborting if
    # newestPackages is empty.
    lib.optionalAttrs (lib.length newestPackages > 0) (versionedDerivations // defaultDerivation);
in
extension