From f797c6df0225760012974d77cde946cd0feac739 Mon Sep 17 00:00:00 2001
From: Mark Canlas <github@htmlism.com>
Date: Sat, 15 Aug 2020 16:23:36 -0400
Subject: [PATCH] allow assignment

---
 .../com/htmlism/mos6502/dsl/package.scala     | 15 +++++++++++---
 .../com/htmlism/mos6502/dsl/DslSpec.scala     | 20 ++++++++++---------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/main/scala/com/htmlism/mos6502/dsl/package.scala b/src/main/scala/com/htmlism/mos6502/dsl/package.scala
index 6fa4263..6c16cca 100644
--- a/src/main/scala/com/htmlism/mos6502/dsl/package.scala
+++ b/src/main/scala/com/htmlism/mos6502/dsl/package.scala
@@ -11,20 +11,29 @@ package object dsl {
       .toDoc
   }
 
-  def group(s: String)(f: DefinitionGroupContext => Unit)(implicit ctx: AsmDocumentContext): Unit = {
+  def group[A](s: String)(f: DefinitionGroupContext => A)(implicit ctx: AsmDocumentContext): A = {
     val g: DefinitionGroupContext =
       new DefinitionGroupContext
 
-    f(g)
+    val ret =
+      f(g)
 
     ctx
       .push(g.toGroup(s))
+
+    ret
   }
 
-  def define[A <: Address : Operand](name: String, x: A)(implicit ctx: DefinitionGroupContext): Unit =
+  def define[A <: Address : Operand](name: String, x: A)(implicit ctx: DefinitionGroupContext): Definition[A] = {
+    val definition =
+      Definition(name, x)
+
     ctx
       .push(Definition(name, x))
 
+    definition
+  }
+
   implicit class AddressOps(n: Int) {
     def z: ZeroAddress =
       ZeroAddress(n)
diff --git a/src/test/scala/com/htmlism/mos6502/dsl/DslSpec.scala b/src/test/scala/com/htmlism/mos6502/dsl/DslSpec.scala
index 85b6ae2..00d3c50 100644
--- a/src/test/scala/com/htmlism/mos6502/dsl/DslSpec.scala
+++ b/src/test/scala/com/htmlism/mos6502/dsl/DslSpec.scala
@@ -9,22 +9,24 @@ class DslSpec extends AnyFlatSpec with should.Matchers {
     val doc =
       asmDoc { implicit ctx =>
         group("snake things") { implicit g =>
-          define("snakeBodyStart", 0x12.z)
-          define("snakeDirection", 0x02.z)
-          define("snakeLength", 0x03.z)
+          (define("snakeBodyStart", 0x12.z),
+            define("snakeDirection", 0x02.z),
+            define("snakeLength", 0x03.z))
         }
 
         group("ASCII values of keys controlling the snake") { implicit g =>
-          define("ASCII_w", 0x77.z)
-          define("ASCII_a", 0x61.z)
-          define("ASCII_s", 0x73.z)
-          define("ASCII_d", 0x64.z)
+          (define("ASCII_w", 0x77.z),
+            define("ASCII_a", 0x61.z),
+            define("ASCII_s", 0x73.z),
+            define("ASCII_d", 0x64.z))
         }
 
         group("System variables") { implicit g =>
-          define("sysRandom", 0xfe.z)
-          define("sysLastKey", 0xff.z)
+          (define("sysRandom", 0xfe.z),
+            define("sysLastKey", 0xff.z))
         }
+
+        ()
       }
 
     doc shouldEqual AsmDocument(List(