Merge pull request #16 from thefrontside/experimental-tail-only-optim…

…ized

Unify stack reduction

mod.ts

@@ -1,190 +1,135 @@
 // deno-lint-ignore-file no-explicit-any
 
-interface Thunk {
-  method: "next" | "throw";
-  iterator: Iterator<Control, unknown, unknown>;
-  value?: unknown | Error;
+export interface Computation<T> {
+  [Symbol.iterator](): Iterator<Control, T, any>;
 }
 
-interface Stack {
-  reducing: boolean;
-  push(...thunks: Thunk[]): number;
-  pop(): Thunk | undefined;
-  value?: any;
+export interface Continuation<T = unknown, R = unknown> {
+  (value: T extends void ? void : T): Computation<R>;
 }
 
-export interface Computation<T = any> {
-  [Symbol.iterator](): Iterator<Control, T, any>;
-}
+export type Result<T> = { ok: true; value: T } | { ok?: false; error: Error };
 
-export interface Continuation<T = any, R = any> {
-  (value: T): R;
-  tail(value: T): void;
-}
+export type Control =
+  | {
+    type: "shift";
+    block(k: Continuation, reject: Continuation<Error>): Computation<unknown>;
+  }
+  | {
+    type: "reset";
+    block(): Computation<unknown>;
+  }
+  | {
+    type: "resume";
+    result: Result<unknown>;
+    iter: Iterator<Control, unknown>;
+  };
 
-export function* reset<T>(block: () => Computation): Computation<T> {
+export function* reset<T>(block: () => Computation<any>): Computation<T> {
   return yield { type: "reset", block };
 }
 
-export function* shift<T>(
-  block: (resolve: Continuation<T>, reject: Continuation<Error>) => Computation,
+export function* shift<T, R = unknown>(
+  block: (
+    k: Continuation<T, R>,
+    reject: Continuation<Error>,
+  ) => Computation<any>,
 ): Computation<T> {
   return yield { type: "shift", block };
 }
 
-function createStack(): Stack {
-  let list: Thunk[] = [];
-  return {
-    reducing: false,
-    push(...thunks: Thunk[]): number {
-      return list.push(...thunks);
-    },
-    pop(): Thunk | undefined {
-      return list.pop();
-    },
-  };
-}
-
-export function evaluate<T>(iterator: () => Computation): T {
-  let stack = createStack();
-  stack.push({
-    method: "next",
-    iterator: iterator()[Symbol.iterator](),
-  });
-  return reduce(stack);
-}
-
-function reduce<T>(stack: Stack): T {
-  try {
-    stack.reducing = true;
-    for (let current = stack.pop(); current; current = stack.pop()) {
-      try {
-        let next = getNext(current);
-        stack.value = next.value;
-
-        if (!next.done) {
-          let control = next.value;
-          if (control.type === "reset") {
-            stack.push(
-              {
-                ...current,
-                method: "next",
-                get value() {
-                  return stack.value;
-                },
-              },
-              {
-                method: "next",
-                iterator: control.block()[Symbol.iterator](),
-              },
-            );
-          } else {
-            let thunk = current;
-            let resolve = oneshot((value: unknown) => {
-              stack.push({
-                method: "next",
-                iterator: thunk.iterator,
-                value,
-              });
-              return reduce(stack);
-            });
-            resolve.tail = oneshot((value: unknown) => {
-              stack.push({
-                method: "next",
-                iterator: thunk.iterator,
-                value,
-              });
-              if (!stack.reducing) {
-                reduce(stack);
-              }
-            });
-            let reject = oneshot((error: Error) => {
-              stack.push({
-                method: "throw",
-                iterator: thunk.iterator,
-                value: error,
-              });
-              return reduce(stack);
-            });
-            reject.tail = oneshot((error: Error) => {
-              stack.push({
-                method: "throw",
-                iterator: thunk.iterator,
-                value: error,
-              });
-
-              if (!stack.reducing) {
-                reduce(stack);
-              }
-            });
-
-            stack.push({
-              method: "next",
-              iterator: control.block(resolve, reject)[Symbol.iterator](),
-              value: void 0,
-            });
-          }
+export function evaluate<T>(block: () => Computation<unknown>) {
+  let stack: Iterator<Control, unknown, any>[] = [];
+  let iter = block()[Symbol.iterator]();
+
+  let current = $next(undefined);
+  while (true) {
+    let result = safe(() => current(iter));
+    if (!result.ok) {
+      if (stack.length > 0) {
+        iter = stack.pop()!;
+        current = $throw(result.error);
+      } else {
+        throw result.error;
+      }
+    } else {
+      let next = result.value;
+      if (next.done) {
+        if (stack.length > 0) {
+          iter = stack.pop()!;
+          current = $next(next.value);
+        } else {
+          return next.value as T;
         }
-      } catch (error) {
-        let top = stack.pop();
-        if (top) {
-          stack.push({ ...top, method: "throw", value: error });
+      } else {
+        const control = next.value;
+        if (control.type === "reset") {
+          stack.push(iter);
+          iter = control.block()[Symbol.iterator]();
+        } else if (control.type === "shift") {
+          const continuation = iter;
+
+          let resolve = oneshot((value: unknown) => ({
+            type: "resume",
+            iter: continuation,
+            result: { ok: true, value },
+          }));
+
+          let reject = oneshot((error: Error) => {
+            return { type: "resume", iter: continuation, result: { error } };
+          });
+          iter = control.block(resolve, reject)[Symbol.iterator]();
         } else {
-          throw error;
+          iter = control.iter;
+          let { result } = control;
+          current = result.ok ? $next(result.value) : $throw(result.error);
         }
       }
     }
-  } finally {
-    stack.reducing = false;
   }
-
-  return stack.value;
 }
 
-function getNext(thunk: Thunk) {
-  let { iterator } = thunk;
-  if (thunk.method === "next") {
-    return iterator.next(thunk.value);
-  } else {
-    let value = thunk.value as Error;
-    if (iterator.throw) {
-      return iterator.throw(value);
-    } else {
-      throw value;
-    }
-  }
+function $next(value: unknown) {
+  return (iter: Iterator<Control, unknown, unknown>) => iter.next(value);
 }
 
-function oneshot<T, R>(fn: (t: T) => R): Continuation<T, R> {
-  let continued = false;
-  let failure: { error: unknown };
-  let result: any;
-
-  return ((value) => {
-    if (!continued) {
-      continued = true;
-      try {
-        return (result = fn(value));
-      } catch (error) {
-        failure = { error };
-        throw error;
-      }
-    } else if (failure) {
-      throw failure.error;
+function $throw(error: Error): ReturnType<typeof $next> {
+  return (iter) => {
+    if (iter.throw) {
+      return iter.throw(error);
     } else {
-      return result;
+      throw error;
     }
-  }) as Continuation<T, R>;
+  };
 }
 
-export type K<T = any, R = any> = Continuation<T, R>;
-
-export type Control =
-  | {
-    type: "shift";
-    block(resolve: Continuation, reject: Continuation<Error>): Computation;
+function safe<T>(fn: () => T): Result<T> {
+  try {
+    return { ok: true, value: fn() };
+  } catch (error) {
+    return { error };
   }
-  | {
-    type: "reset";
-    block(): Computation;
+}
+
+function oneshot<TArg, T>(
+  fn: (arg: TArg) => Control,
+): (arg: TArg) => Computation<T> {
+  let computation: Computation<T> | undefined = undefined;
+
+  return function k(arg: TArg) {
+    if (!computation) {
+      let control = fn(arg);
+      let iterator: Iterator<Control, T> = {
+        next() {
+          iterator.next = () => ({
+            done: true,
+            value: undefined as unknown as T,
+          });
+          return { done: false, value: control };
+        },
+      };
+      computation = { [Symbol.iterator]: () => iterator };
+    }
+    return computation;
   };
+}