Allow setting segments for multiple declarations at once

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Current version
 
+* Added `segment` block statement.
+
 * Added goto.
 
 * Added arrays of elements of size greater than byte.

docs/lang/syntax.md

@@ -155,6 +155,20 @@ TODO
 
 A function can be declared at the top level. For more details, see [Functions](./functions.md)
 
+### Segment blocks
+
+Instead of repeating the segment name in multiple declarations, you can wrap them in a block:
+
+    segment (s) {
+        byte x
+        word y    
+    } 
+    
+is equivalent to:
+
+    segment (s) byte x
+    segment (s) word y
+
 ## `import` statements
 
     import <module>

src/main/scala/millfork/node/Node.scala

@@ -311,6 +311,12 @@ sealed trait DeclarationStatement extends Statement {
   def name: String
 }
 
+sealed trait BankedDeclarationStatement extends DeclarationStatement {
+  def bank: Option[String]
+  def name: String
+  def withChangedBank(bank: String): BankedDeclarationStatement
+}
+
 case class TypeDefinitionStatement(name: String, parent: String) extends DeclarationStatement {
   override def getAllExpressions: List[Expression] = Nil
 }
@@ -325,8 +331,10 @@ case class VariableDeclarationStatement(name: String,
                                         register: Boolean,
                                         initialValue: Option[Expression],
                                         address: Option[Expression],
-                                        alignment: Option[MemoryAlignment]) extends DeclarationStatement {
+                                        alignment: Option[MemoryAlignment]) extends BankedDeclarationStatement {
   override def getAllExpressions: List[Expression] = List(initialValue, address).flatten
+
+  override def withChangedBank(bank: String): BankedDeclarationStatement = copy(bank = Some(bank))
 }
 
 trait ArrayContents extends Node {
@@ -416,8 +424,10 @@ case class ArrayDeclarationStatement(name: String,
                                      address: Option[Expression],
                                      const: Boolean,
                                      elements: Option[ArrayContents],
-                                     alignment: Option[MemoryAlignment]) extends DeclarationStatement {
+                                     alignment: Option[MemoryAlignment]) extends BankedDeclarationStatement {
   override def getAllExpressions: List[Expression] = List(length, address).flatten ++ elements.fold(List[Expression]())(_.getAllExpressions)
+
+  override def withChangedBank(bank: String): BankedDeclarationStatement = copy(bank = Some(bank))
 }
 
 case class ParameterDeclaration(typ: String,
@@ -441,8 +451,10 @@ case class FunctionDeclarationStatement(name: String,
                                         assembly: Boolean,
                                         interrupt: Boolean,
                                         kernalInterrupt: Boolean,
-                                        reentrant: Boolean) extends DeclarationStatement {
+                                        reentrant: Boolean) extends BankedDeclarationStatement {
   override def getAllExpressions: List[Expression] = address.toList ++ statements.getOrElse(Nil).flatMap(_.getAllExpressions)
+
+  override def withChangedBank(bank: String): BankedDeclarationStatement = copy(bank = Some(bank))
 }
 
 sealed trait ExecutableStatement extends Statement

src/main/scala/millfork/parser/MfParser.scala

@@ -102,7 +102,7 @@ abstract class MfParser[T](fileId: String, input: String, currentDirectory: Stri
 
   val atomWithIntel: P[Expression] = P(position() ~ (variableAtom | literalAtomWithIntel | textLiteralAtom)).map{case (p,a) => a.pos(p)}
 
-  val globalVariableDefinition: P[Seq[DeclarationStatement]] = variableDefinition(true)
+  val globalVariableDefinition: P[Seq[BankedDeclarationStatement]] = variableDefinition(true)
   val localVariableDefinition: P[Seq[DeclarationStatement]] = variableDefinition(false)
 
   def singleVariableDefinition: P[(Position, String, Option[Expression], Option[Expression], Option[MemoryAlignment])] = for {
@@ -113,7 +113,7 @@ abstract class MfParser[T](fileId: String, input: String, currentDirectory: Stri
       alignment = None // TODO
     } yield (p, name, addr, initialValue, alignment)
 
-  def variableDefinition(implicitlyGlobal: Boolean): P[Seq[DeclarationStatement]] = for {
+  def variableDefinition(implicitlyGlobal: Boolean): P[Seq[BankedDeclarationStatement]] = for {
     p <- position()
     bank <- bankDeclaration
     flags <- variableFlags ~ HWS
@@ -479,7 +479,7 @@ abstract class MfParser[T](fileId: String, input: String, currentDirectory: Stri
     condition <- "while" ~ !letterOrDigit ~/ HWS ~/ mfExpression(nonStatementLevel, false)
   } yield Seq(DoWhileStatement(body.toList, Nil, condition))
 
-  val functionDefinition: P[Seq[DeclarationStatement]] = for {
+  val functionDefinition: P[Seq[BankedDeclarationStatement]] = for {
     p <- position()
     bank <- bankDeclaration
     flags <- functionFlags ~ HWS
@@ -560,9 +560,22 @@ abstract class MfParser[T](fileId: String, input: String, currentDirectory: Stri
     fields <- compoundTypeFields ~/ Pass
   } yield Seq(UnionDefinitionStatement(name, fields).pos(p))
 
+  val segmentBlock: P[Seq[BankedDeclarationStatement]] = for {
+   bankName <- "segment" ~ AWS ~ "(" ~ AWS ~ identifier ~ AWS ~ ")" ~ AWS ~ "{" ~/ AWS
+   body <- locatableDefinition.rep(sep = EOL)
+   _ <- AWS ~ "}" ~/ Pass
+  } yield {
+    body.flatten.map{ stmt =>
+      if (stmt.bank.isEmpty) stmt.withChangedBank(bankName)
+      else stmt
+    }
+  }
+
+  def locatableDefinition: P[Seq[BankedDeclarationStatement]] = segmentBlock | arrayDefinition | functionDefinition | globalVariableDefinition
+
   val program: Parser[Program] = for {
     _ <- Start ~/ AWS ~/ Pass
-    definitions <- (importStatement | arrayDefinition | aliasDefinition | enumDefinition | structDefinition | unionDefinition | functionDefinition | globalVariableDefinition).rep(sep = EOL)
+    definitions <- (importStatement | aliasDefinition | enumDefinition | structDefinition | unionDefinition | locatableDefinition).rep(sep = EOL)
     _ <- AWS ~ End
   } yield Program(definitions.flatten.toList)
 

src/test/scala/millfork/test/SegmentSuite.scala

@@ -0,0 +1,35 @@
+package millfork.test
+
+import millfork.Cpu
+import millfork.test.emu.EmuUnoptimizedCrossPlatformRun
+import org.scalatest.{FunSuite, Matchers}
+
+/**
+  * @author Karol Stasiak
+  */
+class SegmentSuite extends FunSuite with Matchers {
+  test("Segments") {
+    val source =
+      """
+        | segment(second) byte a1
+        | segment(third) byte a2
+        | segment(second) {
+        |   byte a3
+        |   segment(third) byte a4
+        | }
+        | byte a5
+        | byte output @$c000
+        | void main() {
+        |   output = 0
+        |   if a1.addr.hi & $e0 == $80 { output += 1 }
+        |   if a2.addr.hi & $e0 == $a0 { output += 1 }
+        |   if a3.addr.hi & $e0 == $80 { output += 1 }
+        |   if a4.addr.hi & $e0 == $a0 { output += 1 }
+        |   if a5.addr.hi & $e0 == $00 { output += 1 }
+        | }
+      """.stripMargin
+    EmuUnoptimizedCrossPlatformRun(Cpu.Mos, Cpu.Z80)(source) { m =>
+      m.readByte(0xc000) should equal(source.count(_ == '+'))
+    }
+  }
+}

src/test/scala/millfork/test/emu/EmuPlatform.scala

@@ -18,16 +18,24 @@ object EmuPlatform {
     TextCodec.Ascii,
     Platform.builtInCpuFeatures(cpu),
     CurrentBankFragmentOutput(0, 0xffff),
-    Map("default" -> (if (cpu == Cpu.Intel8086) new UpwardByteAllocator(0x100, 0xb000) else new UpwardByteAllocator(0x200, 0xb000))),
-    Map("default" -> new VariableAllocator(
-      if (CpuFamily.forType(cpu) == CpuFamily.M6502) pointers else Nil,
-      new AfterCodeByteAllocator(0x200, 0xff00))),
+    Map(
+      "default" -> (if (cpu == Cpu.Intel8086) new UpwardByteAllocator(0x100, 0xb000) else new UpwardByteAllocator(0x200, 0xb000)),
+      "second" -> new UpwardByteAllocator(0x8000, 0xa000),
+      "third" -> new UpwardByteAllocator(0xa000, 0xc000)
+    ),
+    Map(
+      "default" -> new VariableAllocator(
+        if (CpuFamily.forType(cpu) == CpuFamily.M6502) pointers else Nil,
+        new AfterCodeByteAllocator(0x200, 0xff00)),
+      "second" -> new VariableAllocator(Nil, new AfterCodeByteAllocator(0x8000, 0xa000)),
+      "third" -> new VariableAllocator(Nil, new AfterCodeByteAllocator(0xa000, 0xc000))
+    ),
     if (CpuFamily.forType(cpu) == CpuFamily.M6502) 4 else 0,
     pointers,
     ".bin",
     false,
     false,
-    Map("default" -> 0),
+    Map("default" -> 0, "second" -> 2, "third" -> 3),
     "default",
     None,
     ViceDebugOutputFormat,