tvl-depot/nix/dependency-analyzer/default.nix

{ lib, depot, pkgs, ... }:

let
  inherit (builtins) unsafeDiscardStringContext appendContext;

  #
  # Utilities
  #

  # Determine all paths a derivation depends on, i.e. input derivations and
  # files imported into the Nix store.
  #
  # Implementation for Nix < 2.6 is quite hacky at the moment.
  #
  # Type: str -> [str]
  #
  # TODO(sterni): clean this up and expose it
  directDrvDeps =
    let
      getDeps =
        if lib.versionAtLeast builtins.nixVersion "2.6"
        then
        # Since https://github.com/NixOS/nix/pull/1643, Nix apparently »preserves
        # string context« through a readFile invocation. This has the side effect
        # that it becomes possible to query the actual references a store path has.
        # Not a 100% sure this is intended, but _very_ convenient for us here.
          drvPath:
          builtins.attrNames (builtins.getContext (builtins.readFile drvPath))
        else
        # For Nix < 2.6 we have to rely on HACK, namely grepping for quoted
        # store path references in the file. In the future this should be
        # replaced by a proper derivation parser.
          drvPath: builtins.concatLists (
            builtins.filter builtins.isList (
              builtins.split
                "\"(${lib.escapeRegex builtins.storeDir}/[[:alnum:]+._?=-]+.drv)\""
                (builtins.readFile drvPath)
            )
          );
    in
    drvPath:
    # if the passed path is not a derivation we can't necessarily get its
    # dependencies, since it may not be representable as a Nix string due to
    # NUL bytes, e.g. compressed patch files imported into the Nix store.
    if builtins.match "^.+\\.drv$" drvPath == null
    then [ ]
    else getDeps drvPath;

  # Maps a list of derivation to the list of corresponding `drvPath`s.
  #
  # Type: [drv] -> [str]
  drvsToPaths = drvs:
    builtins.map (drv: builtins.unsafeDiscardOutputDependency drv.drvPath) drvs;

  #
  # Calculate map of direct derivation dependencies
  #

  # Create the dependency map entry for a given `drvPath` which mainly includes
  # a list of other `drvPath`s it depends on. Additionally we store whether the
  # derivation is `known`, i.e. part of the initial list of derivations we start
  # generating the map from
  #
  # Type: bool -> string -> set
  drvEntry = known: drvPath:
    let
      # key may not refer to a store path, …
      key = unsafeDiscardStringContext drvPath;
      # but we must read from the .drv file.
      path = builtins.unsafeDiscardOutputDependency drvPath;
    in
    {
      inherit key;
      # trick so we can call listToAttrs directly on the result of genericClosure
      name = key;
      value = {
        deps = directDrvDeps path;
        inherit known;
      };
    };

  # Create an attribute set that maps every derivation in the combined
  # dependency closure of the list of input derivation paths to every of their
  # direct dependencies. Additionally every entry will have set their `known`
  # attribute to `true` if it is in the list of input derivation paths.
  #
  # Type: [str] -> set
  plainDrvDepMap = drvPaths:
    builtins.listToAttrs (
      builtins.genericClosure {
        startSet = builtins.map (drvEntry true) drvPaths;
        operator = { value, ... }: builtins.map (drvEntry false) value.deps;
      }
    );

  #
  # Calculate closest known dependencies in the dependency map
  #

  inherit (depot.nix.stateMonad)
    after
    bind
    for_
    get
    getAttr
    run
    setAttr
    pure
    ;

  # This is an action in stateMonad which expects the (initial) state to have
  # been produced by `plainDrvDepMap`. Given a `drvPath`, it calculates a
  # `knownDeps` list which holds the `drvPath`s of the closest derivation marked
  # as `known` along every edge. This list is inserted into the dependency map
  # for `drvPath` and every other derivation in its dependecy closure (unless
  # the information was already present). This means that the known dependency
  # information for a derivation never has to be recalculated, as long as they
  # are part of the same stateful computation.
  #
  # The upshot is that after calling `insertKnownDeps drvPath`,
  # `fmap (builtins.getAttr "knownDeps") (getAttr drvPath)` will always succeed.
  #
  # Type: str -> stateMonad drvDepMap null
  insertKnownDeps = drvPathWithContext:
    let
      # We no longer need to read from the store, so context is irrelevant, but
      # we need to check for attr names which requires the absence of context.
      drvPath = unsafeDiscardStringContext drvPathWithContext;
    in
    bind get (initDepMap:
      # Get the dependency map's state before we've done anything to obtain the
      # entry we'll be manipulating later as well as its dependencies.
      let
        entryPoint = initDepMap.${drvPath};

        # We don't need to recurse if our direct dependencies either have their
        # knownDeps list already populated or are known dependencies themselves.
        depsPrecalculated =
          builtins.partition
            (dep:
              initDepMap.${dep}.known
              || initDepMap.${dep} ? knownDeps
            )
            entryPoint.deps;

        # If a direct dependency is known, it goes right to our known dependency
        # list. If it is unknown, we can copy its knownDeps list into our own.
        initiallyKnownDeps =
          builtins.concatLists (
            builtins.map
              (dep:
                if initDepMap.${dep}.known
                then [ dep ]
                else initDepMap.${dep}.knownDeps
              )
              depsPrecalculated.right
          );
      in

      # If the information was already calculated before, we can exit right away
      if entryPoint ? knownDeps
      then pure null
      else
        after
          # For all unknown direct dependencies which don't have a `knownDeps`
          # list, we call ourselves recursively to populate it. Since this is
          # done sequentially in the state monad, we avoid recalculating the
          # list for the same derivation multiple times.
          (for_
            depsPrecalculated.wrong
            insertKnownDeps)
          # After this we can obtain the updated dependency map which will have
          # a `knownDeps` list for all our direct dependencies and update the
          # entry for the input `drvPath`.
          (bind
            get
            (populatedDepMap:
              (setAttr drvPath (entryPoint // {
                knownDeps =
                  lib.unique (
                    initiallyKnownDeps
                      ++ builtins.concatLists (
                      builtins.map
                        (dep: populatedDepMap.${dep}.knownDeps)
                        depsPrecalculated.wrong
                    )
                  );
              }))))
    );

  # This function puts it all together and is exposed via `__functor`.
  #
  # For a list of `drvPath`s, calculate an attribute set which maps every
  # `drvPath` to a set of the following form:
  #
  #     {
  #       known = true /* if it is in the list of input derivation paths */;
  #       deps = [
  #         /* list of derivation paths it depends on directly */
  #       ];
  #       knownDeps = [
  #         /* list of the closest derivation paths marked as known this
  #            derivation depends on.
  #         */
  #       ];
  #     }
  knownDrvDepMap = knownDrvPaths:
    run
      (plainDrvDepMap knownDrvPaths)
      (after
        (for_
          knownDrvPaths
          insertKnownDeps)
        get);

  #
  # Other things based on knownDrvDepMap
  #

  # Create a SVG visualizing `knownDrvDepMap`. Nodes are identified by derivation
  # name, so multiple entries can be collapsed if they have the same name.
  #
  # Type: [drv] -> drv
  knownDependencyGraph = name: drvs:
    let
      justName = drvPath:
        builtins.substring
          (builtins.stringLength builtins.storeDir + 1 + 32 + 1)
          (builtins.stringLength drvPath)
          (unsafeDiscardStringContext drvPath);

      gv = pkgs.writeText "${name}-dependency-analysis.gv" ''
        digraph depot {
        ${
          (lib.concatStringsSep "\n"
          (lib.mapAttrsToList (name: value:
            if !value.known then ""
            else lib.concatMapStringsSep "\n"
              (knownDep: "  \"${justName name}\" -> \"${justName knownDep}\"")
              value.knownDeps
          )
          (depot.nix.dependency-analyzer (
            drvsToPaths drvs
          ))))
        }
        }
      '';
    in

    pkgs.runCommand "${name}-dependency-analysis.svg"
      {
        nativeBuildInputs = [
          pkgs.buildPackages.graphviz
        ];
      }
      "dot -Tsvg < ${gv} > $out";
in

{
  __functor = _: knownDrvDepMap;

  inherit knownDependencyGraph plainDrvDepMap drvsToPaths;
}