lakefile.lean

import Lake
open Lake DSL

def leanArgs := #["-Dlinter.unusedVariables=false", "-DautoImplicit=false"]

package GroundZero { moreGlobalServerArgs := leanArgs, moreLeanArgs := leanArgs }

@[default_target]
lean_lib GroundZero

set_option linter.unusedVariables false
set_option autoImplicit false

section
  def hasExtension (e : String) : FilePath → Bool :=
  λ fp => fp.extension = some e

  def isTracked (fp : FilePath) : IO Bool :=
  do {
    let retval ← IO.Process.run { cmd := "git", args := #["ls-files", fp.toString] };
    return retval.length > 0
  }

  def modOfPath (fp : FilePath) :=
  ".".intercalate (fp.withExtension "").components

  def modsOfDir (fp : FilePath) :=
  do {
    let fs := (← fp.walkDir).qsort (·.toString < ·.toString)
              |>.filter (hasExtension "lean");
    return (← fs.filterM isTracked).map modOfPath
  }
end

def scriptNotice := "Automatically generated by `lake script run updateIndex`, do not edit it manually."

script updateIndex (argv) do {
  let mods ← modsOfDir "GroundZero";

  IO.FS.withFile "GroundZero.lean" IO.FS.Mode.write fun fd =>
  do { fd.putStrLn s!"-- {scriptNotice}"; for mod in mods do fd.putStrLn s!"import {mod}" };

  return 0
}

section
  universe u v

  open Lean IO

  -- Taken from https://github.com/leanprover-community/mathlib4/blob/master/Mathlib/Util/Imports.lean.
  def importsOf (env : Environment) (n : Name) : Array Name :=
  if n = env.header.mainModule then
    env.header.imports.map Import.module
  else match env.getModuleIdx? n with
  | some idx => env.header.moduleData[idx.toNat]!.imports.map Import.module |>.erase `Init
  | none     => #[]

  def removePrefix (n : Name) :=
  n.replacePrefix `GroundZero Name.anonymous

  def ppName (n : Name) := "/".intercalate (removePrefix n |>.components.map toString)

  def importGraph (E : Environment) : IO (Array (String × String)) := do {
    let mut edges := @Array.empty (String × String);

    for mod₁ in importsOf E `GroundZero do {
      for mod₂ in importsOf E mod₁ do
        if mod₂.getRoot = `GroundZero then
          edges := edges.push (ppName mod₂, ppName mod₁);
    }

    return edges.qsort (λ w₁ w₂ => if w₁.1 = w₂.1 then w₁.2 < w₂.2 else w₁.1 < w₂.1);
  }

  def writeGraph (E : Environment) (outname : String) : IO Unit := do {
    let edges ← importGraph E;

    let child ← Process.spawn {
      cmd := "dot", args := #["-Tsvg", "-o", outname],
      stdout := .inherit, stderr := .inherit, stdin := .piped
    };

    let fd := child.stdin;

    fd.putStrLn "digraph dependencyMap {";
    fd.putStrLn "  splines=polyline";
    for (A, B) in edges do
      fd.putStrLn s!"  \"{A}\" -> \"{B}\""
    fd.putStrLn "}";

    let (_, child) ← child.takeStdin;
    discard child.wait
  }
end

script updateDependencyMap (argv) do {
  Lean.searchPathRef.set (← getWorkspace).augmentedLeanPath;
  let E ← Lean.importModules #[`GroundZero] {} (trustLevel := 1024);
  writeGraph E (argv.getD 0 "/dev/stdout"); return 0
}