From 52e74a4f73cc4325db797da4b77322c50596b18b Mon Sep 17 00:00:00 2001
From: tjjfvi <t6@t6.fyi>
Date: Thu, 22 Jun 2023 16:40:06 -0400
Subject: [PATCH 1/2] flatten cleanup

---
 examples/ink/interact.eg.ts |  4 +-
 fluent/ConnectionRune.ts    | 14 +++++-
 rune/MetaRune.ts            | 41 +++++++++++------
 rune/Rune.ts                | 91 ++++++++++++++++++++++++-------------
 rune/ValueRune.ts           | 34 +++++++-------
 5 files changed, 117 insertions(+), 67 deletions(-)

diff --git a/examples/ink/interact.eg.ts b/examples/ink/interact.eg.ts
index 80cc5fcb9..9cbeb6e85 100644
--- a/examples/ink/interact.eg.ts
+++ b/examples/ink/interact.eg.ts
@@ -8,7 +8,7 @@
 
 import { contractsDev } from "@capi/contracts-dev"
 import { assert } from "asserts"
-import { $, createDevUsers, hex } from "capi"
+import { $, createDevUsers, globalRunner, hex } from "capi"
 import { signature } from "capi/patterns/signature/polkadot"
 import { InkMetadataRune } from "capi/patterns/unstable/ink"
 
@@ -69,6 +69,8 @@ $.assert(
 )
 console.log(events)
 
+console.log([...globalRunner.memo].map(([k, v]) => v._source === k))
+
 /// Retrieve the final state.
 const finalState = await state.run()
 console.log("Alexa final balance:", finalState)
diff --git a/fluent/ConnectionRune.ts b/fluent/ConnectionRune.ts
index 72e7f8810..8fd81c0d5 100644
--- a/fluent/ConnectionRune.ts
+++ b/fluent/ConnectionRune.ts
@@ -3,6 +3,7 @@ import { Connection, ConnectionError, RpcSubscriptionMessage, ServerError } from
 import { is, MetaRune, Run, Rune, RunicArgs, Runner, RunStream } from "../rune/mod.ts"
 
 class RunConnection extends Run<Connection, never> {
+  controller = new AbortController()
   constructor(
     ctx: Runner,
     readonly initConnection: (signal: AbortSignal) => Connection | Promise<Connection>,
@@ -12,7 +13,11 @@ class RunConnection extends Run<Connection, never> {
 
   connection?: Connection
   async _evaluate(): Promise<Connection> {
-    return this.connection ??= await this.initConnection(this.signal)
+    return this.connection ??= await this.initConnection(this.controller.signal)
+  }
+
+  override cleanup(): void {
+    this.controller.abort()
   }
 }
 
@@ -61,6 +66,7 @@ export class ConnectionRune<U> extends Rune<Connection, U> {
 }
 
 class RunRpcSubscription extends RunStream<RpcSubscriptionMessage> {
+  controller = new AbortController()
   constructor(
     ctx: Runner,
     connection: Connection,
@@ -74,7 +80,11 @@ class RunRpcSubscription extends RunStream<RpcSubscriptionMessage> {
       unsubscribeMethod,
       params,
       (value) => this.push(value),
-      this.signal,
+      this.controller.signal,
     )
   }
+
+  override cleanup(): void {
+    this.controller.abort()
+  }
 }
diff --git a/rune/MetaRune.ts b/rune/MetaRune.ts
index 74ef1d21c..2759be2b6 100644
--- a/rune/MetaRune.ts
+++ b/rune/MetaRune.ts
@@ -46,11 +46,11 @@ export class OrthoRune<T, U1, U2> extends Rune<Run<T, U1>, U2> {
 class RunFlat<T, U1, U2> extends Run<T, U1 | U2> {
   child
   constructor(
-    readonly runner: Runner,
+    runner: Runner,
     child: Rune<Rune<T, U1>, U2>,
   ) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   lastChildReceipt = new Receipt()
@@ -62,10 +62,12 @@ class RunFlat<T, U1, U2> extends Run<T, U1 | U2> {
     const rune = await this.child.evaluate(time, receipt)
     if (!receipt.ready) return null!
     if (receipt.novel) {
-      // TODO: prime before aborting?
-      this.innerController.abort()
-      this.innerController = new AbortController()
-      this.currentInner = this.runner.prime(rune, this.innerController.signal)
+      // TODO: prime before dereferencing?
+      if (this.currentInner) {
+        this.currentInner.dereference()
+      }
+      this.currentInner = this.runner.prime(rune)
+      this.currentInner.reference()
     }
     const _receipt = new Receipt()
     try {
@@ -84,7 +86,9 @@ class RunFlat<T, U1, U2> extends Run<T, U1 | U2> {
   }
 
   override cleanup(): void {
-    this.innerController.abort()
+    if (this.currentInner) {
+      this.dependencies.push(this.currentInner)
+    }
     super.cleanup()
   }
 }
@@ -112,13 +116,20 @@ class OrthoRunner extends Runner {
     this.parent.memo.set(rune._prime, run)
     return run
   }
+
+  override onCleanup(run: Run<unknown, unknown>): void {
+    super.onCleanup(run)
+    if (run._sources.length) {
+      this.parent.onCleanup(run)
+    }
+  }
 }
 
 class RunFlatSingular<T, U1, U2> extends Run<T, U1 | U2> {
   child
-  constructor(readonly runner: Runner, child: Rune<Run<T, U1>, U2>) {
+  constructor(runner: Runner, child: Rune<Run<T, U1>, U2>) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
@@ -129,22 +140,24 @@ class RunFlatSingular<T, U1, U2> extends Run<T, U1 | U2> {
 
 class RunAsOrtho<T, U1, U2> extends Run<Run<T, U1>, U2> {
   child
-  constructor(readonly runner: Runner, child: Rune<Rune<T, U1>, U2>) {
+  constructor(runner: Runner, child: Rune<Rune<T, U1>, U2>) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
     const orthoRunner = new OrthoRunner(this.runner, time)
     const rune = await this.child.evaluate(time, receipt)
-    return orthoRunner.prime(rune, undefined)
+    return orthoRunner.prime(rune)
   }
 }
 
 class RunWrapOrtho<T, U> extends Run<T, U> {
-  constructor(readonly runner: OrthoRunner, readonly child: Run<T, U>) {
+  declare runner: OrthoRunner
+
+  constructor(runner: OrthoRunner, readonly child: Run<T, U>) {
     super(runner)
-    child.reference(this.signal)
+    this.useRun(child)
   }
 
   _evaluate(): Promise<T> {
diff --git a/rune/Rune.ts b/rune/Rune.ts
index 66df043a6..14232d421 100644
--- a/rune/Rune.ts
+++ b/rune/Rune.ts
@@ -15,24 +15,27 @@ export abstract class Runner {
   protected abstract _prime<T, U>(rune: Rune<T, U>): Run<T, U>
 
   memo = new Map<(runner: Runner) => Run<any, any>, Run<any, any>>()
-  prime<T, U>(rune: _.Rune<T, U>, signal: AbortSignal | undefined): Run<T, U> {
+  prime<T, U>(rune: Rune<T, U>): Run<T, U> {
     const run = getOrInit(this.memo, rune._prime, () => {
       const old = this._currentTrace
       this._currentTrace = rune._trace
       const run = this._prime(rune)
       this._currentTrace = old
-      run.signal.addEventListener("abort", () => {
-        this.memo.delete(rune._prime)
-      })
+      run._sources.push(rune._prime)
       return run
     })
-    if (signal) run.reference(signal)
     return run
   }
 
-  getPrimed<T, U>(rune: _.Rune<T, U>): Run<T, U> | undefined {
+  getPrimed<T, U>(rune: Rune<T, U>): Run<T, U> | undefined {
     return this.memo.get(rune._prime)
   }
+
+  onCleanup(run: Run<unknown, unknown>) {
+    for (const source of run._sources) {
+      this.memo.delete(source)
+    }
+  }
 }
 
 class RootRunner extends Runner {
@@ -44,7 +47,7 @@ class RootRunner extends Runner {
   }
 }
 
-const globalRunner = new RootRunner()
+export const globalRunner = new RootRunner()
 
 declare const _T: unique symbol
 declare const _U: unique symbol
@@ -84,8 +87,8 @@ export class Rune<out T, out U = never> {
   }
 
   async *iter(runner: Runner = globalRunner) {
-    const abortController = new AbortController()
-    const primed = runner.prime(this, abortController.signal)
+    const primed = runner.prime(this)
+    primed.reference()
     let time = runner.timeline.current
     try {
       while (time !== Infinity) {
@@ -98,7 +101,7 @@ export class Rune<out T, out U = never> {
         time = receipt.nextTime
       }
     } finally {
-      abortController.abort()
+      primed.dereference()
     }
   }
 
@@ -172,7 +175,7 @@ export class Rune<out T, out U = never> {
     )
   }
 
-  static asyncIter<T>(fn: (signal: AbortSignal) => AsyncIterable<T>): Rune.ValueRune<T, never> {
+  static asyncIter<T>(fn: () => AsyncIterable<T>): Rune.ValueRune<T, never> {
     return Rune.ValueRune.new(RunAsyncIter, fn)
   }
 
@@ -203,8 +206,8 @@ export class Rune<out T, out U = never> {
 
   static pin<T, U>(rune: Rune<T, U>, pinned: Rune<unknown, unknown>): Rune<T, U> {
     return new Rune((runner) => {
-      const run = runner.prime(rune, undefined)
-      runner.prime(pinned, run.signal)
+      const run = runner.prime(rune)
+      run.use(pinned)
       return run
     })
   }
@@ -216,27 +219,56 @@ export abstract class Run<T, U> {
   order: number
   timeline
 
-  abortController = new AbortController()
-  signal = this.abortController.signal
-  constructor(runner: Runner) {
-    this.signal.addEventListener("abort", () => this.cleanup())
+  constructor(readonly runner: Runner) {
     this.trace = runner._currentTrace
       ?? new Trace(`execution of the ${new.target.name} instantiated`)
     this.order = runner.order
     this.timeline = runner.timeline
   }
 
+  dependencies: Run<unknown, unknown>[] = []
+
+  use<T, U>(rune: Rune<T, U>): Run<T, U> {
+    const run = this.runner.prime(rune)
+    this.useRun(run)
+    return run
+  }
+
+  useRun(run: Run<unknown, unknown>) {
+    run.reference()
+    this.dependencies.push(run)
+  }
+
   referenceCount = 0
-  reference(signal: AbortSignal) {
+  alive = true
+
+  reference() {
     if (!this.alive) throw new Error("cannot reference a dead rune")
     this.referenceCount++
-    signal.addEventListener("abort", () => {
-      if (!--this.referenceCount) {
-        this.abortController.abort()
+  }
+
+  _sources: Array<(runner: Runner) => Run<any, any>> = []
+  dereference(cleanupBatches?: Run<unknown, unknown>[][]) {
+    if (!--this.referenceCount) {
+      this.alive = false
+      this.cleanup()
+      this.runner.onCleanup(this)
+      if (cleanupBatches) {
+        cleanupBatches.push(this.dependencies)
+      } else {
+        const cleanupBatches = [this.dependencies]
+        while (cleanupBatches.length) {
+          const batch = cleanupBatches.pop()!
+          for (const run of batch) {
+            run.dereference(cleanupBatches)
+          }
+        }
       }
-    })
+    }
   }
 
+  cleanup() {}
+
   _currentTime = -1
   _currentPromise: Promise<T> = null!
   _currentReceipt = new Receipt()
@@ -258,11 +290,6 @@ export abstract class Run<T, U> {
     }
   }
   abstract _evaluate(time: number, receipt: Receipt): Promise<T>
-
-  alive = true
-  cleanup() {
-    this.alive = false
-  }
 }
 
 class RunConstant<T> extends Run<T, never> {
@@ -279,7 +306,7 @@ class RunLs<T, U> extends Run<T[], U> {
   children
   constructor(runner: Runner, children: Rune<T, U>[]) {
     super(runner)
-    this.children = children.map((child) => runner.prime(child, this.signal))
+    this.children = children.map((child) => this.use(child))
   }
 
   _evaluate(time: number, receipt: Receipt) {
@@ -335,10 +362,10 @@ export abstract class RunStream<T> extends Run<T, never> {
 }
 
 class RunAsyncIter<T> extends RunStream<T> {
-  constructor(runner: Runner, fn: (signal: AbortSignal) => AsyncIterable<T>) {
+  constructor(runner: Runner, fn: () => AsyncIterable<T>) {
     super(runner)
     ;(async () => {
-      for await (const value of fn(this.signal)) {
+      for await (const value of fn()) {
         this.push(value)
       }
       this.finish()
@@ -356,7 +383,7 @@ class RunBubbleUnhandled<T, U> extends Run<T, never> {
   child
   constructor(runner: Runner, child: Rune<T, U>, readonly symbol: symbol) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
@@ -375,7 +402,7 @@ class RunCaptureUnhandled<T, U1, U2> extends Run<T, U1 | U2> {
   child
   constructor(runner: Runner, child: Rune<T, U1>, readonly symbol: symbol) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
diff --git a/rune/ValueRune.ts b/rune/ValueRune.ts
index 600aa563a..3eb431c0c 100644
--- a/rune/ValueRune.ts
+++ b/rune/ValueRune.ts
@@ -182,10 +182,8 @@ class RunMatch<M, T, U> extends Run<T, U> {
     conditions: [(x: M) => boolean, ValueRune<T, U>][],
   ) {
     super(runner)
-    this.value = runner.prime(child, this.signal)
-    this.conditions = conditions.map(([cond, val]) =>
-      [cond, runner.prime(val, this.signal)] as const
-    )
+    this.value = this.use(child)
+    this.conditions = conditions.map(([cond, val]) => [cond, this.use(val)] as const)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
@@ -207,7 +205,7 @@ class RunMap<T1, U, T2> extends Run<T2, U> {
     readonly fn: (value: T1) => T2 | Promise<T2>,
   ) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   lastValue: T2 = null!
@@ -231,8 +229,8 @@ class RunHandle<T, T2 extends T, T3, U, U2> extends Run<Exclude<T, T2> | T3, U |
     alt: Rune<T3, U2>,
   ) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
-    this.alt = runner.prime(alt, this.signal)
+    this.child = this.use(child)
+    this.alt = this.use(alt)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
@@ -251,7 +249,7 @@ class RunUnhandle<T, U> extends Run<T, U | T> {
     readonly guard: Guard<T, T>,
   ) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
@@ -269,7 +267,7 @@ class RunThrows<T, U1, U2> extends Run<T, U1 | U2> {
     readonly guards: Array<Guard<unknown, U2>>,
   ) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
@@ -288,7 +286,7 @@ class RunGetUnhandled<T, U> extends Run<Unhandled<U> | null, never> {
   child
   constructor(runner: Runner, child: Rune<T, U>) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
@@ -314,8 +312,8 @@ class RunRehandle<T1, U1, U2 extends U1, T3, U3> extends Run<T1 | T3, Exclude<U1
     alt: Rune<T3, U3>,
   ) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
-    this.alt = runner.prime(alt, this.signal)
+    this.child = this.use(child)
+    this.alt = this.use(alt)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
@@ -334,7 +332,7 @@ class RunLazy<T, U> extends Run<T, U> {
   child
   constructor(runner: Runner, child: Rune<T, U>) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, _receipt: Receipt): Promise<T> {
@@ -346,7 +344,7 @@ class RunFilter<T, U> extends Run<T, U> {
   child
   constructor(runner: Runner, child: Rune<T, U>, readonly fn: (value: T) => boolean) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   first = true
@@ -375,7 +373,7 @@ class RunFinal<T, U> extends Run<T, U> {
   child
   constructor(runner: Runner, child: Rune<T, U>) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, receipt: Receipt): Promise<T> {
@@ -395,7 +393,7 @@ class RunReduce<T1, U, T2> extends Run<T2, U> {
     readonly fn: (last: T2, value: T1) => T2 | Promise<T2>,
   ) {
     super(runner)
-    this.child = runner.prime(child, this.signal)
+    this.child = this.use(child)
   }
 
   async _evaluate(time: number, receipt: Receipt) {
@@ -415,8 +413,8 @@ class RunChain<T1, U1, T2, U2> extends Run<T2, U1 | U2> {
     second: Rune<T2, U2>,
   ) {
     super(runner)
-    this.first = runner.prime(first, this.signal)
-    this.second = runner.prime(second, this.signal)
+    this.first = this.use(first)
+    this.second = this.use(second)
   }
 
   lastValue: T2 = null!

From af53f68048b02772310473809bd791bb46426d69 Mon Sep 17 00:00:00 2001
From: tjjfvi <t6@t6.fyi>
Date: Fri, 23 Jun 2023 08:39:10 -0400
Subject: [PATCH 2/2] revert old testing code

---
 examples/ink/interact.eg.ts | 4 +---
 rune/Rune.ts                | 2 +-
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/examples/ink/interact.eg.ts b/examples/ink/interact.eg.ts
index 9cbeb6e85..80cc5fcb9 100644
--- a/examples/ink/interact.eg.ts
+++ b/examples/ink/interact.eg.ts
@@ -8,7 +8,7 @@
 
 import { contractsDev } from "@capi/contracts-dev"
 import { assert } from "asserts"
-import { $, createDevUsers, globalRunner, hex } from "capi"
+import { $, createDevUsers, hex } from "capi"
 import { signature } from "capi/patterns/signature/polkadot"
 import { InkMetadataRune } from "capi/patterns/unstable/ink"
 
@@ -69,8 +69,6 @@ $.assert(
 )
 console.log(events)
 
-console.log([...globalRunner.memo].map(([k, v]) => v._source === k))
-
 /// Retrieve the final state.
 const finalState = await state.run()
 console.log("Alexa final balance:", finalState)
diff --git a/rune/Rune.ts b/rune/Rune.ts
index 14232d421..9669b7481 100644
--- a/rune/Rune.ts
+++ b/rune/Rune.ts
@@ -47,7 +47,7 @@ class RootRunner extends Runner {
   }
 }
 
-export const globalRunner = new RootRunner()
+const globalRunner = new RootRunner()
 
 declare const _T: unique symbol
 declare const _U: unique symbol