Add Haskell code

.gitignore

@@ -0,0 +1,4 @@
+/dist/
+/dist-newstyle/
+/cabal.project.local
+/cabal.project.local~*

app/Main.hs

@@ -0,0 +1,88 @@
+{-# LANGUAGE ApplicativeDo #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE QuasiQuotes #-}
+
+module Main (main) where
+
+import Data.List.NonEmpty (NonEmpty)
+import Numeric.Natural (Natural)
+
+import qualified Algebra.Graph.Export.Dot
+import qualified Data.List.NonEmpty as NonEmpty
+import qualified Nix.Graph
+import qualified Options.Applicative as Options
+
+data Options = Options
+  { paths :: NonEmpty FilePath
+  , excludeBuilt :: Bool
+  , excludeCached :: Bool
+  , maxFiles :: Natural
+  }
+
+parseOptions :: Options.Parser Options
+parseOptions = do
+  paths <-
+    fmap NonEmpty.fromList $
+      Options.some $ Options.strArgument (Options.metavar "PATH")
+
+  excludeBuilt <-
+    Options.switch . mconcat $
+      [ Options.long "exclude-built"
+      , Options.help . mconcat $
+          [ "Reduce size of graph by excluding derivations which have already"
+          , "been built, and have their output path in the local Nix store."
+          ]
+      , Options.hidden
+      ]
+
+  excludeCached <-
+    Options.switch . mconcat $
+      [ Options.long "exclude-cached"
+      , Options.help . mconcat $
+          [ "Reduce size of graph by excluding derivations which do not need to"
+          , "be built (because they have already been built locally) or can be"
+          , "substituted from another cache. Implies `--exclude-built`."
+          ]
+      , Options.hidden
+      ]
+
+  maxFiles <-
+    Options.option Options.auto . mconcat $
+      [ Options.long "max-files"
+      , Options.metavar "COUNT"
+      , Options.help "Limit number of open files"
+      , Options.value 100
+      , Options.showDefault
+      , Options.hidden
+      ]
+
+  pure Options{paths, excludeBuilt, excludeCached, maxFiles}
+
+getOptions :: IO Options
+getOptions = do
+  let parserPrefs =
+        Options.prefs . mconcat $
+          [ Options.showHelpOnError
+          , Options.multiSuffix ".."
+          ]
+  let parserInfo = Options.info (Options.helper <*> parseOptions) mempty
+  Options.customExecParser parserPrefs parserInfo
+
+main :: IO ()
+main = do
+  Options{paths, excludeBuilt, excludeCached, maxFiles} <- getOptions
+
+  let exclude
+        | excludeCached = Nix.Graph.ExcludeCached
+        | excludeBuilt = Nix.Graph.ExcludeBuilt
+        | otherwise = Nix.Graph.ExcludeNothing
+
+  let config =
+        Nix.Graph.Config
+          { Nix.Graph.exclude
+          , Nix.Graph.maxFiles
+          }
+
+  adjacencyMap <- Nix.Graph.build config (NonEmpty.toList paths)
+
+  putStrLn (Algebra.Graph.Export.Dot.exportAsIs adjacencyMap)

nix-graph.cabal

@@ -0,0 +1,66 @@
+cabal-version: 2.2
+
+name:        nix-graph
+version:     1.0.0.0
+synopsis:    Reify the Nix build graph into a Haskell graph data structure
+description: Reify the Nix build graph into a Haskell graph data structure
+category:    Nix, Graphs
+author:      Arista Networks <[email protected]>
+maintainer:  Arista Networks <[email protected]>
+homepage:    https://github.com/awakesecurity/nix-graph
+license:     BSD-3-Clause
+copyright:   2021 Arista Networks
+tested-with: GHC == 8.8.4, GHC == 8.10.4
+
+license-file: LICENSE
+extra-source-files:
+  CHANGELOG.md
+  LICENSE
+  README.md
+
+common common
+  build-depends: base >= 4.13 && < 5.0
+  default-language: Haskell2010
+  ghc-options:
+    -Wall
+    -Wcompat
+    -Werror=incomplete-record-updates
+    -Werror=incomplete-uni-patterns
+    -Werror=missing-fields
+    -Werror=partial-fields
+    -Widentities
+    -Wmissing-home-modules
+    -Wredundant-constraints
+    -fshow-warning-groups
+
+common executable
+  ghc-options:
+    -threaded
+    -rtsopts
+    -with-rtsopts=-N
+
+library
+  import: common
+  hs-source-dirs: src
+  build-depends:
+    , algebraic-graphs >= 0.5    && < 0.6
+    , attoparsec       >= 0.13.1 && < 0.15
+    , containers
+    , hashable
+    , nix-derivation   >= 1.1.0  && < 2.0
+    , stm              >= 2.4.2
+    , text
+    , ttrie            >= 0.1.2  && < 0.2
+    , unliftio         >= 0.2.6  && < 0.3
+  exposed-modules:
+    Nix.Graph
+    Nix.Graph.Internal
+
+executable nix-graph
+  import: common, executable
+  hs-source-dirs: app
+  main-is: Main.hs
+  build-depends:
+    , algebraic-graphs
+    , nix-graph
+    , optparse-applicative >= 0.15.0.0

src/Nix/Graph.hs

@@ -0,0 +1,7 @@
+module Nix.Graph (
+  Config (..),
+  Exclude (..),
+  build,
+) where
+
+import Nix.Graph.Internal

src/Nix/Graph/Internal.hs

@@ -0,0 +1,198 @@
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+module Nix.Graph.Internal where
+
+import Algebra.Graph.AdjacencyMap (AdjacencyMap)
+import Control.Applicative ((<|>))
+import Control.Concurrent.STM.TSem (TSem)
+import Control.Monad (unless, when)
+import Control.Monad.IO.Class (MonadIO (..))
+import Data.Attoparsec.Text ((<?>))
+import Data.Hashable (Hashable)
+import Data.Set (Set)
+import GHC.Generics (Generic)
+import Numeric.Natural (Natural)
+import System.Exit (ExitCode (..))
+
+import qualified Algebra.Graph.AdjacencyMap as AdjacencyMap
+import qualified Control.Concurrent.STM.Map as STM.Map
+import qualified Control.Concurrent.STM.TSem as TSem
+import qualified Data.Attoparsec.Text as Attoparsec
+import qualified Data.Map as Map
+import qualified Data.Set as Set
+import qualified Data.Text as Text
+import qualified Data.Text.IO as Text.IO
+import qualified Nix.Derivation
+import qualified UnliftIO.Async as Async
+import qualified UnliftIO.Directory as Directory
+import qualified UnliftIO.Exception as Exception
+import qualified UnliftIO.IO as IO
+import qualified UnliftIO.Process as Process
+import qualified UnliftIO.STM as STM
+
+data Derivation = Derivation
+  { derivationPath :: FilePath
+  , derivationInputDrvs :: [FilePath]
+  , derivationBuilt :: Bool
+  }
+  deriving stock (Eq, Ord, Generic)
+  deriving anyclass (Hashable)
+
+readDerivation :: (MonadIO m, MonadFail m) => TSem -> FilePath -> m Derivation
+readDerivation tSem derivationPath = do
+  let acquire = STM.atomically $ TSem.waitTSem tSem
+  let release = STM.atomically $ TSem.signalTSem tSem
+
+  drv <- liftIO . Exception.bracket_ acquire release $
+    IO.withFile derivationPath IO.ReadMode $ \handle -> do
+      fileContents <- Text.IO.hGetContents handle
+      case Attoparsec.parseOnly Nix.Derivation.parseDerivation fileContents of
+        Left err -> fail ("Failed to parse derivation: " <> err)
+        Right drv -> pure drv
+
+  outputPath <-
+    case Map.lookup "out" (Nix.Derivation.outputs drv) of
+      Nothing -> fail "Failed to lookup output path"
+      Just output -> pure (Nix.Derivation.path output)
+
+  derivationBuilt <- Directory.doesPathExist outputPath
+
+  let derivationInputDrvs = Map.keys (Nix.Derivation.inputDrvs drv)
+
+  pure Derivation{derivationPath, derivationBuilt, derivationInputDrvs}
+
+buildAdjacencyMap ::
+  MonadIO m =>
+  Eq k =>
+  Hashable k =>
+  (k -> IO (Set k)) ->
+  [k] ->
+  m [(k, Set k)]
+buildAdjacencyMap getNeighbors roots = liftIO $ do
+  stmMap <- STM.atomically STM.Map.empty
+
+  let go key = do
+        isMember <- STM.atomically $ do
+          isMember <- STM.Map.member key stmMap
+          unless isMember $ STM.Map.insert key Set.empty stmMap
+          pure isMember
+        unless isMember $ do
+          neighbors <- getNeighbors key
+          STM.atomically $ STM.Map.insert key neighbors stmMap
+          Async.mapConcurrently_ go neighbors
+
+  Async.mapConcurrently_ go roots
+
+  STM.Map.unsafeToList stmMap
+
+filterUnbuilt :: (MonadIO m, MonadFail m) => [FilePath] -> m (Set FilePath)
+filterUnbuilt derivationPaths = do
+  (exitCode, _nixStdOut, nixStdErr) <-
+    Process.readProcessWithExitCode
+      "nix-store"
+      ( [ "--realize"
+        , "--dry-run"
+        ]
+          <> derivationPaths
+      )
+      ""
+
+  when (exitCode /= ExitSuccess) $ do
+    fail ("Failed to run 'nix-store --realize --dry-run " <> unwords derivationPaths <> "':\n" <> nixStdErr)
+
+  case Attoparsec.parseOnly willBeBuilt (Text.pack nixStdErr) of
+    Left parseError ->
+      fail ("Failed to parse output from 'nix-store --realize --dry-run ...':\n" <> parseError)
+    Right derivationPathsToBuild -> do
+      pure (Set.fromList derivationPathsToBuild)
+  where
+    willBeBuilt :: Attoparsec.Parser [FilePath]
+    willBeBuilt = Attoparsec.option [] $ do
+      willBeBuiltHeading
+      Attoparsec.many' nixStorePath
+
+    -- Slightly different headings depending on Nix version
+    --
+    -- Nix 2: https://github.com/NixOS/nix/blob/2.3.10/src/libmain/shared.cc#L45-L71
+    -- (uses the same strings from Nix 2.0 to 2.3.10)
+
+    --
+    -- Nix 3: https://github.com/NixOS/nix/blob/8e758d4/src/libmain/shared.cc#L48-L86
+    -- (latest as of 2020-02-25)
+
+    willBeBuiltHeading :: Attoparsec.Parser ()
+    willBeBuiltHeading = do
+      let nix2 = "these derivations will be built:"
+      let nix3 =
+            "this derivation will be built:"
+              <|> ("these " *> Attoparsec.decimal @Int *> " derivations will be built:")
+      _ <- nix2 <|> nix3
+      Attoparsec.endOfLine
+
+    nixStorePath :: Attoparsec.Parser FilePath
+    nixStorePath = (<?> "/nix/store path") $ do
+      _ <- "  "
+      nixStore <- "/nix/store"
+      rest <- Attoparsec.takeTill Attoparsec.isEndOfLine
+      Attoparsec.endOfLine
+      pure (Text.unpack (nixStore <> rest))
+
+data Config = Config
+  { exclude :: Exclude
+  , maxFiles :: Natural
+  }
+
+data Exclude
+  = ExcludeNothing
+  | ExcludeBuilt
+  | ExcludeCached
+  deriving stock (Eq)
+
+-- | Build graph of dependencies
+build ::
+  MonadIO m =>
+  -- | Configure how the graph is built
+  Config ->
+  -- | Derivations to build graph from
+  [FilePath] ->
+  m (AdjacencyMap FilePath)
+build _ [] = pure (AdjacencyMap.empty)
+build Config{exclude, maxFiles} roots = liftIO $ do
+  tSem <- STM.atomically $ TSem.newTSem (toInteger maxFiles)
+
+  process :: [FilePath] -> IO [Derivation] <- do
+    case exclude of
+      ExcludeCached -> do
+        unbuiltSet <- filterUnbuilt roots
+
+        pure \paths -> do
+          let unbuiltPaths = filter (`Set.member` unbuiltSet) paths
+
+          Async.mapConcurrently (readDerivation tSem) unbuiltPaths
+      ExcludeBuilt -> do
+        pure \paths -> do
+          derivations <- Async.mapConcurrently (readDerivation tSem) paths
+
+          pure (filter (not . derivationBuilt) derivations)
+      ExcludeNothing -> do
+        pure \paths -> do
+          Async.mapConcurrently (readDerivation tSem) paths
+
+  let getInputDrvs :: Derivation -> IO (Set Derivation)
+      getInputDrvs derivation = do
+        fmap Set.fromList (process (derivationInputDrvs derivation))
+
+  rootDrvs <- process roots
+
+  adjacencySets <- buildAdjacencyMap getInputDrvs rootDrvs
+
+  let adjacencyMap = AdjacencyMap.fromAdjacencySets adjacencySets
+
+  pure (AdjacencyMap.gmap derivationPath adjacencyMap)