Zig-like "defer" to clean up stuff when leaving the scope of the current routine.

codeCore/src/prog8/code/ast/AstPrinter.kt

@@ -154,6 +154,7 @@ fun printAst(root: PtNode, skipLibraries: Boolean, output: (text: String) -> Uni
                 else
                     "->"
             }
+            is PtDefer -> "<defer>"
         }
     }
 

codeCore/src/prog8/code/ast/AstStatements.kt

@@ -179,3 +179,6 @@ class PtWhenChoice(val isElse: Boolean, position: Position) : PtNode(position) {
     val statements: PtNodeGroup
         get() = children[1] as PtNodeGroup
 }
+
+
+class PtDefer(position: Position): PtNode(position), IPtStatementContainer

codeCore/src/prog8/code/core/BuiltinFunctions.kt

@@ -129,6 +129,7 @@ val BuiltinFunctions: Map<String, FSignature> = mapOf(
     "callfar"   to FSignature(false, listOf(FParam("bank", arrayOf(DataType.UBYTE)), FParam("address", arrayOf(DataType.UWORD)), FParam("arg", arrayOf(DataType.UWORD))), DataType.UWORD),
     "callfar2"  to FSignature(false, listOf(FParam("bank", arrayOf(DataType.UBYTE)), FParam("address", arrayOf(DataType.UWORD)), FParam("argA", arrayOf(DataType.UBYTE)), FParam("argX", arrayOf(DataType.UBYTE)), FParam("argY", arrayOf(DataType.UBYTE)), FParam("argC", arrayOf(DataType.BOOL))), DataType.UWORD),
     "call"      to FSignature(false, listOf(FParam("address", arrayOf(DataType.UWORD))), DataType.UWORD),
+    "invoke_defer"  to FSignature(false, emptyList(), null),
 )
 
 val InplaceModifyingBuiltinFunctions = setOf(

codeCore/src/prog8/code/core/SourceCode.kt

@@ -9,6 +9,7 @@ import kotlin.io.path.readText
 
 
 const val internedStringsModuleName = "prog8_interned_strings"
+const val deferLabel = "prog8_defer_statements"
 
 
 /**

codeGenCpu6502/src/prog8/codegen/cpu6502/AsmGen.kt

@@ -607,6 +607,7 @@ class AsmGen6502Internal (
             is PtVariable, is PtConstant, is PtMemMapped -> { /* do nothing; variables are handled elsewhere */ }
             is PtBlock -> throw AssemblyError("block should have been handled elsewhere")
             is PtNodeGroup -> stmt.children.forEach { translate(it) }
+            is PtDefer -> translate(stmt)
             is PtNop -> {}
             else -> throw AssemblyError("missing asm translation for $stmt")
         }
@@ -1092,6 +1093,15 @@ $repeatLabel""")
         }
     }
 
+    private fun translate(defer: PtDefer) {
+        val sub = defer.definingSub()!!
+        out("${sub.name}.$deferLabel")
+        for(stmt in defer.children) {
+            translate(stmt)
+        }
+        out("  rts")
+    }
+
     internal fun signExtendAYlsb(valueDt: DataType) {
         // sign extend signed byte in A to full word in AY
         when(valueDt) {

codeGenCpu6502/src/prog8/codegen/cpu6502/BuiltinFunctionsAsmGen.kt

@@ -69,12 +69,18 @@ internal class BuiltinFunctionsAsmGen(private val program: PtProgram,
             "prog8_lib_stringcompare" -> funcStringCompare(fcall, resultRegister)
             "prog8_lib_square_byte" -> funcSquare(fcall, DataType.UBYTE, resultRegister)
             "prog8_lib_square_word" -> funcSquare(fcall, DataType.UWORD, resultRegister)
+            "invoke_defer" -> funcInvokeDefer(fcall)
             else -> throw AssemblyError("missing asmgen for builtin func ${fcall.name}")
         }
 
         return BuiltinFunctions.getValue(fcall.name).returnType
     }
 
+    private fun funcInvokeDefer(call: PtBuiltinFunctionCall) {
+        val sub = call.definingSub()!!
+        asmgen.out("  jsr  ${sub.name}.$deferLabel")
+    }
+
     private fun funcSquare(fcall: PtBuiltinFunctionCall, resultType: DataType, resultRegister: RegisterOrPair?) {
         // square of word value is faster with dedicated routine, square of byte just use the regular multiplication routine.
         when (resultType) {

codeGenIntermediate/src/prog8/codegen/intermediate/BuiltinFuncGen.kt

@@ -44,10 +44,19 @@ internal class BuiltinFuncGen(private val codeGen: IRCodeGen, private val exprGe
             "prog8_lib_stringcompare" -> funcStringCompare(call)
             "prog8_lib_square_byte" -> funcSquare(call, IRDataType.BYTE)
             "prog8_lib_square_word" -> funcSquare(call, IRDataType.WORD)
+            "invoke_defer" -> funcInvokeDefer(call)
             else -> throw AssemblyError("missing builtinfunc for ${call.name}")
         }
     }
 
+    private fun funcInvokeDefer(call: PtBuiltinFunctionCall): ExpressionCodeResult {
+        val sub = call.definingSub()!!
+        val result = mutableListOf<IRCodeChunkBase>()
+        addInstr(result, IRInstruction(Opcode.CALL, labelSymbol = "${sub.name}.$deferLabel",
+            fcallArgs = FunctionCallArgs(emptyList(), emptyList())), null)
+        return ExpressionCodeResult(result, IRDataType.BYTE, -1, -1)
+    }
+
     private fun funcSquare(call: PtBuiltinFunctionCall, resultType: IRDataType): ExpressionCodeResult {
         val result = mutableListOf<IRCodeChunkBase>()
         val valueTr = exprGen.translateExpression(call.args[0])

codeGenIntermediate/src/prog8/codegen/intermediate/IRCodeGen.kt

@@ -237,6 +237,7 @@ class IRCodeGen(
                 listOf(chunk)
             }
             is PtConditionalBranch -> translate(node)
+            is PtDefer -> translate(node)
             is PtInlineAssembly -> listOf(IRInlineAsmChunk(null, node.assembly, node.isIR, null))
             is PtIncludeBinary -> listOf(IRInlineBinaryChunk(null, readBinaryData(node), null))
             is PtAddressOf,
@@ -274,6 +275,16 @@ class IRCodeGen(
             .map { it.toUByte() }
     }
 
+    private fun translate(defer: PtDefer): IRCodeChunks {
+        val result = mutableListOf<IRCodeChunkBase>()
+        for(stmt in defer.children) {
+            result += translateNode(stmt)
+        }
+        addInstr(result, IRInstruction(Opcode.RETURN), null)
+        val sub = defer.definingSub()!!
+        return labelFirstChunk(result, "${sub.name}.$deferLabel")
+    }
+
     private fun translate(branch: PtConditionalBranch): IRCodeChunks {
         val result = mutableListOf<IRCodeChunkBase>()
 

compiler/src/prog8/compiler/Compiler.kt

@@ -131,6 +131,10 @@ fun compileProgram(args: CompilerArguments): CompilationResult? {
                 val intermediateAst = IntermediateAstMaker(program, args.errors).transform()
                 val stMaker = SymbolTableMaker(intermediateAst, compilationOptions)
                 val symbolTable = stMaker.make()
+
+                postprocessIntermediateAst(intermediateAst, symbolTable, args.errors)
+                args.errors.report()
+
                 if(compilationOptions.optimize) {
                     optimizeIntermediateAst(intermediateAst, compilationOptions, symbolTable, args.errors)
                     args.errors.report()

compiler/src/prog8/compiler/astprocessing/IntermediateAstMaker.kt

@@ -37,6 +37,7 @@ class IntermediateAstMaker(private val program: Program, private val errors: IEr
     private fun transformStatement(statement: Statement): PtNode {
         return when (statement) {
             is AnonymousScope -> throw FatalAstException("AnonymousScopes should have been flattened")
+            is Defer -> transform(statement)
             is ChainedAssignment -> throw FatalAstException("ChainedAssignment should have been flattened")
             is Assignment -> transform(statement)
             is Block -> transform(statement)
@@ -88,6 +89,14 @@ class IntermediateAstMaker(private val program: Program, private val errors: IEr
         }
     }
 
+    private fun transform(srcDefer: Defer): PtDefer {
+        val defer = PtDefer(srcDefer.position)
+        srcDefer.scope.statements.forEach {
+            defer.add(transformStatement(it))
+        }
+        return defer
+    }
+
     private fun transform(srcAssign: Assignment): PtNode {
         if(srcAssign.isAugmentable) {
             require(srcAssign.target.multi==null)

compiler/src/prog8/compiler/astprocessing/IntermediateAstPostprocess.kt

@@ -0,0 +1,85 @@
+package prog8.compiler.astprocessing
+
+import prog8.code.SymbolTable
+import prog8.code.ast.*
+import prog8.code.core.DataType
+import prog8.code.core.IErrorReporter
+
+internal fun postprocessIntermediateAst(program: PtProgram, st: SymbolTable, errors: IErrorReporter) {
+    coalesceDefers(program)
+    integrateDefers(program, st)
+}
+
+
+private fun coalesceDefers(program: PtProgram) {
+    val defersPerSub = mutableMapOf<PtSub, MutableList<PtDefer>>().withDefault { mutableListOf() }
+
+    walkAst(program) { node, _ ->
+        if(node is PtDefer) {
+            val scope = node.definingSub()!!
+            val defers = defersPerSub.getValue(scope)
+            defers.add(node)
+            defersPerSub[scope] = defers
+        }
+    }
+
+    for((sub, defers) in defersPerSub) {
+        val coalescedDefer = PtDefer(sub.position)
+        for(defer in defers.reversed()) {
+            for(stmt in defer.children)
+                coalescedDefer.add(stmt)
+            sub.children.remove(defer)
+        }
+
+        if(coalescedDefer.children.isNotEmpty()) {
+            sub.add(coalescedDefer)
+        }
+    }
+}
+
+
+private fun integrateDefers(program: PtProgram, st: SymbolTable) {
+    val exitsToAugment = mutableListOf<PtNode>()
+    val subEndsToAugment = mutableListOf<PtSub>()
+
+    walkAst(program) { node, _ ->
+        when(node) {
+            is PtJump -> {
+                if(node.identifier!=null) {
+                    val stNode = st.lookup(node.identifier!!.name)!!
+                    val targetSub = stNode.astNode.definingSub()
+                    if(targetSub!=node.definingSub())
+                        exitsToAugment.add(node)
+                }
+            }
+            is PtReturn -> exitsToAugment.add(node)
+            is PtSub -> {
+                val lastStmt = node.children.lastOrNull { it !is PtDefer }
+                if(lastStmt != null && lastStmt !is PtReturn && lastStmt !is PtJump)
+                    subEndsToAugment.add(node)
+            }
+            else -> {}
+        }
+    }
+
+    for(exit in exitsToAugment) {
+        val defer = exit.definingSub()!!.children.singleOrNull { it is PtDefer }
+        if(defer != null) {
+            val idx = exit.parent.children.indexOf(exit)
+            val invokedefer = PtBuiltinFunctionCall("invoke_defer", true, false, DataType.UNDEFINED, exit.position)
+            exit.parent.add(idx, invokedefer)
+        }
+    }
+
+    for(sub in subEndsToAugment) {
+        val defer = sub.children.singleOrNull { it is PtDefer }
+        if(defer != null) {
+            val idx = sub.children.indexOfLast { it !is PtDefer }
+            val ret = PtReturn(sub.position)
+            sub.add(idx+1, ret)
+            val invokedefer = PtBuiltinFunctionCall("invoke_defer", true, false, DataType.UNDEFINED, sub.position)
+            sub.add(idx+1, invokedefer)
+        }
+    }
+
+}

compilerAst/src/prog8/ast/AstToSourceTextConverter.kt

@@ -430,6 +430,14 @@ class AstToSourceTextConverter(val output: (text: String) -> Unit, val program:
         outputi("}")
     }
 
+    override fun visit(defer: Defer) {
+        outputln("defer {")
+        scopelevel++
+        outputStatements(defer.scope.statements)
+        scopelevel--
+        outputi("}")
+    }
+
     override fun visit(typecast: TypecastExpression) {
         output("(")
         typecast.expression.accept(this)

compilerAst/src/prog8/ast/antlr/Antlr2Kotlin.kt

@@ -156,6 +156,9 @@ private fun StatementContext.toAst() : Statement {
     val unrollstmt = unrollloop()?.toAst()
     if(unrollstmt!=null) return unrollstmt
 
+    val deferstmt = defer()?.toAst()
+    if(deferstmt!=null) return deferstmt
+
     throw FatalAstException("unprocessed source text (are we missing ast conversion rules for parser elements?): $text")
 }
 
@@ -642,6 +645,17 @@ private fun BreakstmtContext.toAst() = Break(toPosition())
 
 private fun ContinuestmtContext.toAst() = Continue(toPosition())
 
+private fun DeferContext.toAst(): Defer {
+    val block = statement_block()?.toAst()
+    if(block!=null) {
+        val scope = AnonymousScope(block, statement_block()?.toPosition() ?: toPosition())
+        return Defer(scope, toPosition())
+    }
+    val singleStmt = statement()!!.toAst()
+    val scope = AnonymousScope(mutableListOf(singleStmt), statement().toPosition())
+    return Defer(scope, toPosition())
+}
+
 private fun WhileloopContext.toAst(): WhileLoop {
     val condition = expression().toAst()
     val statements = statement_block()?.toAst() ?: mutableListOf(statement().toAst())

compilerAst/src/prog8/ast/statements/AstStatements.kt

@@ -731,6 +731,24 @@ class InlineAssembly(val assembly: String, val isIR: Boolean, override val posit
     }
 }
 
+class Defer(val scope: AnonymousScope, override val position: Position): Statement() {
+    override lateinit var parent: Node
+
+    override fun linkParents(parent: Node) {
+        this.parent = parent
+        scope.linkParents(this)
+    }
+
+    override fun replaceChildNode(node: Node, replacement: Node) {
+        TODO("Not yet implemented")
+    }
+
+    override fun referencesIdentifier(nameInSource: List<String>): Boolean = scope.referencesIdentifier(nameInSource)
+    override fun copy() = Defer(scope.copy(), position)
+    override fun accept(visitor: IAstVisitor) = visitor.visit(this)
+    override fun accept(visitor: AstWalker, parent: Node) = visitor.visit(this, parent)
+}
+
 class AnonymousScope(override var statements: MutableList<Statement>,
                      override val position: Position) : IStatementContainer, Statement() {
     override lateinit var parent: Node

compilerAst/src/prog8/ast/walk/AstWalker.kt

@@ -130,6 +130,7 @@ abstract class AstWalker {
     open fun before(untilLoop: UntilLoop, parent: Node): Iterable<IAstModification> = noModifications
     open fun before(returnStmt: Return, parent: Node): Iterable<IAstModification> = noModifications
     open fun before(scope: AnonymousScope, parent: Node): Iterable<IAstModification> = noModifications
+    open fun before(defer: Defer, parent: Node): Iterable<IAstModification> = noModifications
     open fun before(char: CharLiteral, parent: Node): Iterable<IAstModification> = noModifications
     open fun before(string: StringLiteral, parent: Node): Iterable<IAstModification> = noModifications
     open fun before(subroutine: Subroutine, parent: Node): Iterable<IAstModification> = noModifications
@@ -174,6 +175,7 @@ abstract class AstWalker {
     open fun after(untilLoop: UntilLoop, parent: Node): Iterable<IAstModification> = noModifications
     open fun after(returnStmt: Return, parent: Node): Iterable<IAstModification> = noModifications
     open fun after(scope: AnonymousScope, parent: Node): Iterable<IAstModification> = noModifications
+    open fun after(defer: Defer, parent: Node): Iterable<IAstModification> = noModifications
     open fun after(char: CharLiteral, parent: Node): Iterable<IAstModification> = noModifications
     open fun after(string: StringLiteral, parent: Node): Iterable<IAstModification> = noModifications
     open fun after(subroutine: Subroutine, parent: Node): Iterable<IAstModification> = noModifications
@@ -442,6 +444,12 @@ abstract class AstWalker {
         track(after(scope, parent), scope, parent)
     }
 
+    fun visit(defer: Defer, parent: Node) {
+        track(before(defer, parent), defer, parent)
+        defer.scope.accept(this, defer)
+        track(after(defer, parent), defer, parent)
+    }
+
     fun visit(typecast: TypecastExpression, parent: Node) {
         track(before(typecast, parent), typecast, parent)
         typecast.expression.accept(this, typecast)

compilerAst/src/prog8/ast/walk/IAstVisitor.kt

@@ -159,6 +159,10 @@ interface IAstVisitor {
         scope.statements.forEach { it.accept(this) }
     }
 
+    fun visit(defer: Defer) {
+        defer.scope.accept(this)
+    }
+
     fun visit(typecast: TypecastExpression) {
         typecast.expression.accept(this)
     }

docs/source/todo.rst

@@ -1,6 +1,11 @@
 TODO
 ====
 
+- fix defer that a return <expression> is evaluated first (and saved), then the defer is called, then a simple return <value> is done
+- unit test for defer
+- describe defer in the manual
+
+
 Improve register load order in subroutine call args assignments:
 in certain situations, the "wrong" order of evaluation of function call arguments is done which results
 in overwriting registers that already got their value, which requires a lot of stack juggling (especially on plain 6502 cpu!)
@@ -46,7 +51,6 @@ Future Things and Ideas
   But all library code written in asm uses .proc already..... (textual search/replace when writing the actual asm?)
   Once new codegen is written that is based on the IR, this point is mostly moot anyway as that will have its own dead code removal on the IR level.
 - Zig-like try-based error handling where the V flag could indicate error condition? and/or BRK to jump into monitor on failure? (has to set BRK vector for that) But the V flag is also set on certain normal instructions
-- Zig-like defer to clean up stuff when leaving the scope. But as we don't have closures, it's more limited. Still useful?
 
 
 Libraries:

examples/test.p8

@@ -4,13 +4,38 @@
 
 main {
     sub start() {
-        ubyte v0
-        ubyte v1
-        ubyte v2
-        ubyte v3
-        v0 = v1 = v2 = 99
-        for v3 in 10 to 20 {
-            cx16.r0L++
+        ubyte x = testdefer()
+        txt.print("result from call=")
+        txt.print_ub(x)
+        txt.nl()
+    }
+
+    sub testdefer() -> ubyte {
+        ubyte var = 22
+
+        defer txt.print("defer1\n")
+        defer {
+            txt.print("defer2, var=")
+            txt.print_ub(var)
+            txt.nl()
         }
+
+        if var==22 {
+            var = 88
+            return var
+        }
+        else {
+            var++
+            txt.print("var=")
+            txt.print_ub(var)
+            txt.nl()
+            return 255
+        }
+
+
+    }
+
+    sub other() {
+        cx16.r0++
     }
 }

parser/antlr/Prog8ANTLR.g4

@@ -109,6 +109,7 @@ statement :
 	| breakstmt
 	| continuestmt
 	| labeldef
+	| defer
 	;
 
 
@@ -126,6 +127,7 @@ subroutinedeclaration :
 	| romsubroutine
     ;
 
+defer: 'defer' (statement | statement_block) ;
 
 labeldef :  identifier ':'  ;
 

syntax-files/IDEA/Prog8.xml

@@ -14,7 +14,7 @@
     <keywords keywords="&amp;;-&gt;;@;and;as;asmsub;break;clobbers;continue;do;downto;else;false;for;goto;if;if_cc;if_cs;if_eq;if_mi;if_ne;if_neg;if_nz;if_pl;if_pos;if_vc;if_vs;if_z;in;inline;not;or;repeat;return;romsub;step;sub;to;true;unroll;until;when;while;xor;~" ignore_case="false" />
     <keywords2 keywords="%address;%asm;%asmbinary;%asminclude;%breakpoint;%encoding;%import;%ir;%launcher;%option;%output;%zeropage;%zpallowed;%zpreserved;atascii:;cp437:;default:;iso16:;iso5:;iso:;kata:;petscii:;sc:" />
     <keywords3 keywords="@nozp;@requirezp;@shared;@split;@zp;bool;byte;const;float;str;ubyte;uword;void;word" />
-    <keywords4 keywords="abs;call;callfar;callfar2;clamp;cmp;divmod;len;lsb;max;memory;min;mkword;msb;peek;peekf;peekw;poke;pokef;pokew;rol;rol2;ror;ror2;rrestore;rrestorex;rsave;rsavex;setlsb;setmsb;sgn;sizeof;sqrt" />
+    <keywords4 keywords="abs;call;callfar;callfar2;clamp;cmp;defer;divmod;len;lsb;max;memory;min;mkword;msb;peek;peekf;peekw;poke;pokef;pokew;rol;rol2;ror;ror2;rrestore;rrestorex;rsave;rsavex;setlsb;setmsb;sgn;sizeof;sqrt" />
   </highlighting>
   <extensionMap>
     <mapping ext="p8" />

syntax-files/NotepadPlusPlus/Prog8.xml

@@ -27,7 +27,7 @@
             <Keywords name="Keywords1">void const&#x000D;&#x000A;str&#x000D;&#x000A;byte ubyte bool&#x000D;&#x000A;word uword&#x000D;&#x000A;float&#x000D;&#x000A;zp shared split requirezp nozp</Keywords>
             <Keywords name="Keywords2">%address&#x000D;&#x000A;%asm&#x000D;&#x000A;%ir&#x000D;&#x000A;%asmbinary&#x000D;&#x000A;%asminclude&#x000D;&#x000A;%breakpoint&#x000D;&#x000A;%encoding&#x000D;&#x000A;%import&#x000D;&#x000A;%launcher&#x000D;&#x000A;%option&#x000D;&#x000A;%output&#x000D;&#x000A;%zeropage&#x000D;&#x000A;%zpreserved&#x000D;&#x000A;%zpallowed</Keywords>
             <Keywords name="Keywords3">inline sub asmsub romsub&#x000D;&#x000A;clobbers&#x000D;&#x000A;asm&#x000D;&#x000A;if&#x000D;&#x000A;when else&#x000D;&#x000A;if_cc if_cs if_eq if_mi if_neg if_nz if_pl if_pos if_vc if_vs if_z&#x000D;&#x000A;for in step do while repeat unroll&#x000D;&#x000A;break continue return goto</Keywords>
-            <Keywords name="Keywords4">abs call callfar callfar2 clamp cmp divmod len lsb lsl lsr memory mkword min max msb peek peekw peekf poke pokew pokef rsave rsavex rrestore rrestorex rnd rndw rol rol2 ror ror2 setlsb setmsb sgn sizeof sqrtw</Keywords>
+            <Keywords name="Keywords4">abs call callfar callfar2 clamp cmp defer divmod len lsb lsl lsr memory mkword min max msb peek peekw peekf poke pokew pokef rsave rsavex rrestore rrestorex rnd rndw rol rol2 ror ror2 setlsb setmsb sgn sizeof sqrtw</Keywords>
             <Keywords name="Keywords5">true false&#x000D;&#x000A;not and or xor&#x000D;&#x000A;as to downto |&gt;</Keywords>
             <Keywords name="Keywords6"></Keywords>
             <Keywords name="Keywords7"></Keywords>

syntax-files/Vim/prog8_builtins.vim

@@ -15,7 +15,7 @@ syn keyword prog8BuiltInFunc len
 " Miscellaneous functions
 syn keyword prog8BuiltInFunc cmp divmod lsb msb mkword min max peek peekw peekf poke pokew pokef rsave rsavex rrestore rrestorex
 syn keyword prog8BuiltInFunc rol rol2 ror ror2 sizeof setlsb setmsb
-syn keyword prog8BuiltInFunc memory call callfar callfar2 clamp
+syn keyword prog8BuiltInFunc memory call callfar callfar2 clamp defer
 
 
 " c64/floats.p8