refactor viz state updates (#9630)

* refactor viz state updates

* onclick

tinygrad/viz/index.html

@@ -274,16 +274,19 @@
   var ret = [];
   var cache = {};
   var kernels = null;
-  var currentUOp = 0;
-  var currentKernel = -1;
-  var currentRewrite = 0;
-  var expandKernel = true;
   const evtSources = [];
+  const state = {currentKernel:-1, currentUOp:0, currentRewrite:0, expandKernel:false};
+  function setState(ns) {
+    Object.assign(state, ns);
+    main();
+  }
   async function main() {
+    // NOTE: this is readonly now
+    const { currentKernel, currentUOp, currentRewrite, expandKernel } = state;
     // ***** LHS kernels list
     if (kernels == null) {
       kernels = await (await fetch("/kernels")).json();
-      currentKernel = -1;
+      setState({ currentKernel:-1 });
     }
     const kernelList = document.querySelector(".kernel-list");
     kernelList.innerHTML = "";
@@ -305,24 +308,13 @@
         rwUl.style.display = i === currentKernel && expandKernel ? "block" : "none";
         rwUl.onclick = (e) => {
           e.stopPropagation();
-          currentUOp = j;
-          currentKernel = i;
-          currentRewrite = 0;
-          main();
+          setState({ currentUOp:j, currentKernel:i, currentRewrite:0 });
         }
         kernelUl.appendChild(rwUl)
       })
       p.onclick = () => {
-        if (i === currentKernel) {
-          expandKernel = !expandKernel;
-          main();
-          return;
-        }
-        currentKernel = i;
-        currentUOp = 0;
-        currentRewrite = 0;
-        expandKernel = true;
-        main();
+        if (i === currentKernel) return setState({ expandKernel:!expandKernel });
+        setState({ expandKernel:true, currentKernel:i, currentUOp:0, currentRewrite:0 });
       }
       kernelList.appendChild(kernelUl);
     });
@@ -400,10 +392,7 @@
             metadata.appendChild(Object.assign(document.createElement("pre"), { innerHTML: `<code>${diffHtml}</code>`, className: "wrap" }));
           }
         }
-        gUl.addEventListener("click", () => {
-          currentRewrite = i;
-          main();
-        });
+        gUl.addEventListener("click", () => setState({ currentRewrite: i }));
       }
     } else {
       metadata.appendChild(Object.assign(document.createElement("p"), { textContent: `No rewrites in ${toPath(kernel.loc)}.` }));
@@ -444,60 +433,43 @@
 
   // **** keyboard shortcuts
   document.addEventListener("keydown", async function(event) {
+    const { currentKernel, currentUOp, currentRewrite, expandKernel } = state;
     // up and down change the UOp or kernel from the list
     if (!expandKernel) {
       if (event.key == "ArrowUp") {
         event.preventDefault()
-        currentUOp = 0;
-        currentRewrite = 0;
-        currentKernel = Math.max(0, currentKernel-1)
-       return main()
+        return setState({ currentUOp:0, currentRewrite:0, currentKernel:Math.max(0, currentKernel-1) });
       }
       if (event.key == "ArrowDown") {
         event.preventDefault()
-        currentUOp = 0;
-        currentRewrite = 0;
-        currentKernel = Math.min(kernels.length-1, currentKernel+1);
-       return main()
+        return setState({ currentUOp:0, currentRewrite:0, currentKernel:Math.max(kernels.length-1, currentKernel+1) });
       }
     }
     if (event.key == "Enter") {
       event.preventDefault()
-      if (currentKernel === -1) {
-        currentKernel = 0;
-        expandKernel = true;
+      if (state.currentKernel === -1) {
+        return setState({ currentKernel:0, expandKernel:true });
       }
-      else {
-        expandKernel = !expandKernel;
-      }
-      currentUOp = 0;
-      currentRewrite = 0;
-      main();
+      return setState({ currentUOp:0, currentRewrite:0, expandKernel:!expandKernel });
     }
     if (event.key == "ArrowUp") {
       event.preventDefault()
-      currentRewrite = 0;
-      currentUOp = Math.max(0, currentUOp-1)
-      main()
+      return setState({ currentRewrite:0, currentUOp:Math.max(0, currentUOp-1) });
     }
     if (event.key == "ArrowDown") {
       event.preventDefault()
-      currentRewrite = 0;
       const totalUOps = kernels[currentKernel][1].length-1;
-      currentUOp = Math.min(totalUOps, currentUOp+1)
-      main()
+      return setState({ currentRewrite:0, currentUOp:Math.min(totalUOps, currentUOp+1) });
     }
     // left and right go through rewrites in a single UOp
     if (event.key == "ArrowLeft") {
       event.preventDefault()
-      currentRewrite = Math.max(0, currentRewrite-1)
-      main()
+      return setState({ currentRewrite:Math.max(0, currentRewrite-1) });
     }
     if (event.key == "ArrowRight") {
       event.preventDefault()
       const totalRewrites = ret.length-1;
-      currentRewrite = Math.min(totalRewrites, currentRewrite+1)
-      main()
+      return setState({ currentRewrite:Math.min(totalRewrites, currentRewrite+1) });
     }
     if (event.key == " ") {
       document.getElementById("zoom-to-fit-btn").click();