From f32780d8635b8eca60f09e61f04a77f261bc363e Mon Sep 17 00:00:00 2001
From: Mac Malone <tydeu@hatpress.net>
Date: Thu, 20 Jun 2024 21:43:05 -0400
Subject: [PATCH] refactor: lake: more robust trace reading (#4518)

The recent change of the trace format exposed some unexpected issues
with Lake's tracing handling. This aims to fix that.

Lake will now perform a rebuild if the trace file is invalid/unreadable.
However, it will still fall back to modification times if the trace file
is missing. Also, Lake is now backwards compatible with the previous
pure numeric traces (and tolerates the absence of a `log` field in the
JSON trace).
---
 src/lake/Lake/Build/Common.lean    | 66 ++++++++++++++++++++----------
 src/lake/Lake/Build/Trace.lean     |  5 ++-
 src/lake/Lake/Util/JsonObject.lean |  3 ++
 src/lake/tests/trace/Foo.lean      |  0
 src/lake/tests/trace/clean.sh      |  1 +
 src/lake/tests/trace/lakefile.toml |  5 +++
 src/lake/tests/trace/test.sh       | 35 ++++++++++++++++
 7 files changed, 93 insertions(+), 22 deletions(-)
 create mode 100644 src/lake/tests/trace/Foo.lean
 create mode 100755 src/lake/tests/trace/clean.sh
 create mode 100644 src/lake/tests/trace/lakefile.toml
 create mode 100755 src/lake/tests/trace/test.sh

diff --git a/src/lake/Lake/Build/Common.lean b/src/lake/Lake/Build/Common.lean
index 642303db8bbe..a03f4b2f8e2d 100644
--- a/src/lake/Lake/Build/Common.lean
+++ b/src/lake/Lake/Build/Common.lean
@@ -5,6 +5,7 @@ Authors: Mac Malone
 -/
 import Lake.Config.Monad
 import Lake.Build.Actions
+import Lake.Util.JsonObject
 
 /-! # Common Build Tools
 This file defines general utilities that abstract common
@@ -31,15 +32,29 @@ which stores information about a (successful) build.
 structure BuildMetadata where
   depHash : Hash
   log : Log
-  deriving ToJson, FromJson
+  deriving ToJson
+
+def BuildMetadata.ofHash (h : Hash) : BuildMetadata :=
+  {depHash := h, log := {}}
+
+def BuildMetadata.fromJson? (json : Json) : Except String BuildMetadata := do
+  let obj ← JsonObject.fromJson? json
+  let depHash ← obj.get "depHash"
+  let log ← obj.getD "log" {}
+  return {depHash, log}
+
+instance : FromJson BuildMetadata := ⟨BuildMetadata.fromJson?⟩
 
 /-- Read persistent trace data from a file. -/
 def readTraceFile? (path : FilePath) : LogIO (Option BuildMetadata) := OptionT.run do
   match (← IO.FS.readFile path |>.toBaseIO) with
   | .ok contents =>
-    match Json.parse contents >>= fromJson? with
-    | .ok contents => return contents
-    | .error e => logVerbose s!"{path}: invalid trace file: {e}"; failure
+    if let some hash := Hash.ofString? contents.trim then
+      return .ofHash hash
+    else
+      match Json.parse contents >>= fromJson? with
+      | .ok contents => return contents
+      | .error e => logVerbose s!"{path}: invalid trace file: {e}"; failure
   | .error (.noFileOrDirectory ..) => failure
   | .error e => logWarning s!"{path}: read failed: {e}"; failure
 
@@ -86,25 +101,34 @@ then `depTrace` / `oldTrace`. No log will be replayed.
   (depTrace : BuildTrace) (traceFile : FilePath) (build : JobM PUnit)
   (action : JobAction := .build) (oldTrace := depTrace.mtime)
 : JobM Bool := do
-  if let some data ← readTraceFile? traceFile then
-    if (← checkHashUpToDate info depTrace data.depHash oldTrace) then
-      updateAction .replay
-      data.log.replay
-      return true
-  else if (← getIsOldMode) then
-    if (← oldTrace.checkUpToDate info) then
+  if (← traceFile.pathExists) then
+    if let some data ← readTraceFile? traceFile then
+      if (← checkHashUpToDate info depTrace data.depHash oldTrace) then
+        updateAction .replay
+        data.log.replay
+        return true
+      else
+        go
+    else if (← getIsOldMode) && (← oldTrace.checkUpToDate info) then
       return true
-  else if (← depTrace.checkAgainstTime info) then
-    return true
-  if (← getNoBuild) then
-    IO.Process.exit noBuildCode.toUInt8
+    else
+      go
   else
-    updateAction action
-    let iniPos ← getLogPos
-    build -- fatal errors will not produce a trace (or cache their log)
-    let log := (← getLog).takeFrom iniPos
-    writeTraceFile traceFile depTrace log
-    return false
+    if (← depTrace.checkAgainstTime info) then
+      return true
+    else
+      go
+where
+  go := do
+    if (← getNoBuild) then
+      IO.Process.exit noBuildCode.toUInt8
+    else
+      updateAction action
+      let iniPos ← getLogPos
+      build -- fatal errors will not produce a trace (or cache their log)
+      let log := (← getLog).takeFrom iniPos
+      writeTraceFile traceFile depTrace log
+      return false
 
 /--
 Checks whether `info` is up-to-date, and runs `build` to recreate it if not.
diff --git a/src/lake/Lake/Build/Trace.lean b/src/lake/Lake/Build/Trace.lean
index 0eb7bf1d9987..87e2cfe4091c 100644
--- a/src/lake/Lake/Build/Trace.lean
+++ b/src/lake/Lake/Build/Trace.lean
@@ -87,8 +87,11 @@ namespace Hash
 @[inline] def ofNat (n : Nat) :=
   mk n.toUInt64
 
+def ofString? (s : String) : Option Hash :=
+  (inline s.toNat?).map ofNat
+
 def load? (hashFile : FilePath) : BaseIO (Option Hash) :=
-  (·.toNat?.map ofNat) <$> IO.FS.readFile hashFile |>.catchExceptions fun _ => pure none
+  ofString? <$> IO.FS.readFile hashFile |>.catchExceptions fun _ => pure none
 
 def nil : Hash :=
   mk <| 1723 -- same as Name.anonymous
diff --git a/src/lake/Lake/Util/JsonObject.lean b/src/lake/Lake/Util/JsonObject.lean
index 2ae8f8ee5725..0327e8517486 100644
--- a/src/lake/Lake/Util/JsonObject.lean
+++ b/src/lake/Lake/Util/JsonObject.lean
@@ -20,6 +20,9 @@ abbrev JsonObject :=
 
 namespace JsonObject
 
+@[inline] def mk (val : RBNode String (fun _ => Json)) : JsonObject :=
+  val
+
 @[inline] protected def toJson (obj : JsonObject) : Json :=
   .obj obj
 
diff --git a/src/lake/tests/trace/Foo.lean b/src/lake/tests/trace/Foo.lean
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/src/lake/tests/trace/clean.sh b/src/lake/tests/trace/clean.sh
new file mode 100755
index 000000000000..b6e54c681e54
--- /dev/null
+++ b/src/lake/tests/trace/clean.sh
@@ -0,0 +1 @@
+rm -rf .lake lake-manifest.json
diff --git a/src/lake/tests/trace/lakefile.toml b/src/lake/tests/trace/lakefile.toml
new file mode 100644
index 000000000000..0b869b10f7db
--- /dev/null
+++ b/src/lake/tests/trace/lakefile.toml
@@ -0,0 +1,5 @@
+name = "test"
+defaultTargets = ["Foo"]
+
+[[lean_lib]]
+name = "Foo"
diff --git a/src/lake/tests/trace/test.sh b/src/lake/tests/trace/test.sh
new file mode 100755
index 000000000000..bcfab40bf0bb
--- /dev/null
+++ b/src/lake/tests/trace/test.sh
@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+set -euxo pipefail
+
+LAKE=${LAKE:-../../.lake/build/bin/lake}
+
+./clean.sh
+
+# ---
+# Tests aspects of Lake tracing
+# ---
+
+# Tests that a build produces a trace
+test ! -f .lake/build/lib/Foo.trace
+$LAKE build | grep --color "Built Foo"
+test -f .lake/build/lib/Foo.trace
+
+# Tests that a proper trace prevents a rebuild
+$LAKE build --no-build
+
+# Tests that Lake accepts pure numerical traces
+if command -v jq > /dev/null; then # skip if no jq found
+  jq -r '.depHash' .lake/build/lib/Foo.trace > .lake/build/lib/Foo.trace.hash
+  mv .lake/build/lib/Foo.trace.hash .lake/build/lib/Foo.trace
+  $LAKE build --no-build
+fi
+
+# Tests that removal of the trace does not cause a rebuild
+# (if the modification time of the artifact is still newer than the inputs)
+rm .lake/build/lib/Foo.trace
+$LAKE build --no-build
+
+# Tests that an invalid trace does cause a rebuild
+touch .lake/build/lib/Foo.trace
+$LAKE build | grep --color "Built Foo"
+$LAKE build --no-build