From e8ec3f544b40edd308b550cc820652eaa3c360c8 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Mon, 16 Jun 2025 12:26:10 +0300
Subject: [PATCH] viz: helper for keeping state in browser history (#10830)

---
 tinygrad/viz/js/index.js | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index 61b5f1d55e..9f021076ed 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -38,14 +38,7 @@ async function renderDag(graph, additions, recenter=false) {
     const STROKE_WIDTH = 1.4;
     const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g")
       .attr("transform", d => `translate(${d.x},${d.y})`).classed("clickable", d => d.ref != null)
-      .on("click", (_,d) => {
-        if (d.ref != null) {
-          // NOTE: browser does a structured clone, passing a mutable object is safe.
-          history.replaceState(state, "");
-          history.pushState(state, "");
-          setState({ expandSteps: true, currentCtx:d.ref, currentStep:0, currentRewrite:0 });
-        }
-      });
+      .on("click", (_,d) => setCtxWithHistory(d.ref));
     nodes.selectAll("rect").data(d => [d]).join("rect").attr("width", d => d.width).attr("height", d => d.height).attr("fill", d => d.color)
       .attr("x", d => -d.width/2).attr("y", d => -d.height/2).attr("style", d => d.style ?? `stroke:#4a4b57; stroke-width:${STROKE_WIDTH}px;`);
     nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label").attr("transform", d => {
@@ -302,6 +295,16 @@ function setState(ns) {
   // re-render
   main();
 }
+
+// set a new context and keep the old one in browser history
+function setCtxWithHistory(newCtx) {
+  if (newCtx == null) return;
+  // NOTE: browser does a structured clone, passing a mutable object is safe.
+  history.replaceState(state, "");
+  history.pushState(state, "");
+  setState({ expandSteps:true, currentCtx:newCtx, currentStep:0, currentRewrite:0 });
+}
+
 window.addEventListener("popstate", (e) => {
   if (e.state != null) setState(e.state);
 });