Fix passing a value to a resumed generator

src/ark/compiler/index.ts

@@ -172,10 +172,7 @@ export function flatToJs(insts: ArkInsts, file: string | null = null): CodeWithS
       } else if (inst instanceof ArkContinueInst) {
         return sourceNode('continue\n')
       } else if (inst instanceof ArkYieldInst) {
-        return sourceNode([
-          `yield ${inst.argId.description}\n`,
-          letAssign(inst.id, inst.argId.description!),
-        ])
+        return sourceNode(letAssign(inst.id, `yield ${inst.argId.description}`))
       } else if (inst instanceof ArkReturnInst) {
         return sourceNode(`return ${inst.argId.description}\n`)
       } else if (inst instanceof ArkLetCopyInst) {

src/ark/interpreter.ts

@@ -814,7 +814,7 @@ async function evalFlat(outerArk: ArkState): Promise<ArkVal> {
       inst = ark.loopStack[0]
     } else if (inst instanceof ArkYieldInst) {
       const result = mem.get(inst.argId)!
-      ark.inst = inst.next
+      ark.inst = inst
       if (ark.outerState === undefined || ark.outerState.continuation === undefined) {
         throw new ArkRuntimeError(ark, 'yield outside a generator', inst.sourceLoc)
       }
@@ -877,6 +877,11 @@ async function evalFlat(outerArk: ArkState): Promise<ArkVal> {
           callable.state.outerState = ark
           ark = callable.state
           inst = ark.inst
+          // If we're resuming, 'inst' pointed to the ArkYieldInst so we can
+          // set its result, so we need to advance to the next instruction.
+          if (inst instanceof ArkYieldInst) {
+            inst = inst.next
+          }
         }
       } else if (callable instanceof NativeFn) {
         mem.set(inst.id, callable.body(...args))

src/grammar/ursa.ohm

@@ -107,6 +107,7 @@ Ursa {
     | Loop
     | For
     | await Exp       -- await
+    | yield Exp?      -- yield
     | LogicExp
 
   Assignment
@@ -120,7 +121,6 @@ Ursa {
     | continue        -- continue
     | launch Exp      -- launch
     | return Exp?     -- return
-    | yield Exp?      -- yield
     | Exp
 
   Lets = NonemptyListOf<Let, and>

src/ursa/compiler.ts

@@ -525,6 +525,13 @@ semantics.addOperation<ArkExp>('toExp(a)', {
     return addLoc(new ArkAwait(exp.toExp(this.args.a)), this)
   },
 
+  Exp_yield(yield_, exp) {
+    if (!this.args.a.inGenerator) {
+      throw new UrsaCompilerError(yield_.source, 'yield may only be used in a generator')
+    }
+    return addLoc(new ArkYield(maybeVal(this.args.a, exp)), this)
+  },
+
   Statement_break(_break, exp) {
     if (!this.args.a.inLoop) {
       throw new UrsaCompilerError(_break.source, 'break used outside a loop')
@@ -548,14 +555,6 @@ semantics.addOperation<ArkExp>('toExp(a)', {
     }
     return addLoc(new ArkReturn(maybeVal(this.args.a, exp)), this)
   },
-  Statement_yield(yield_, exp) {
-    if (!this.args.a.inGenerator) {
-      throw new UrsaCompilerError(yield_.source, 'yield may only be used in a generator')
-    } else if (this.args.a.inExp) {
-      throw new UrsaCompilerError(yield_.source, 'yield may not be used inside an expression')
-    }
-    return addLoc(new ArkYield(maybeVal(this.args.a, exp)), this)
-  },
 
   Block(_open, seq, _close) {
     return seq.toExp(this.args.a)

src/ursa/examples.test.ts

@@ -63,7 +63,11 @@ ursaTest('yield', 'test/yield', [], `\
 3
 4
 null
-null`)
+null
+0
+1
+5
+5`)
 ursaTest('generator', 'test/generator', [], `\
 0
 2

src/ursa/fmt.ts

@@ -463,6 +463,14 @@ semantics.addOperation<Span>('fmt(a)', {
     )
   },
 
+  Exp_yield(_yield, exp) {
+    return tryFormats(
+      this.args.a,
+      (a) => hSpan(['yield', ...fmtOptional(a, exp)]),
+      [(a) => hSpan(['yield', ...fmtOptional(a, exp)])],
+    )
+  },
+
   Statement_break(_break, exp) {
     return tryFormats(
       this.args.a,
@@ -487,13 +495,6 @@ semantics.addOperation<Span>('fmt(a)', {
       [(a) => hSpan(['return', ...fmtOptional(a, exp)])],
     )
   },
-  Statement_yield(_yield, exp) {
-    return tryFormats(
-      this.args.a,
-      (a) => hSpan(['yield', ...fmtOptional(a, exp)]),
-      [(a) => hSpan(['yield', ...fmtOptional(a, exp)])],
-    )
-  },
 
   Lets(lets) {
     return tryFormats(

test/yield.ursa

@@ -17,3 +17,15 @@ print(h())
 print(h())
 print(h())
 print(h())
+
+let totalizer = gen() {
+    let i = 0
+    loop {
+        i := i + (yield i)
+    }
+}
+let t = totalizer()
+print(t(0))
+print(t(1))
+print(t(4))
+print(t(0))