now also support using defer inside if statements

compiler/src/prog8/compiler/astprocessing/AstChecker.kt

@@ -1184,9 +1184,6 @@ internal class AstChecker(private val program: Program,
         defer.scope.accept(s)
         if(s.count>0)
             errors.err("defer cannot contain jumps or returns", defer.position)
-
-        if(defer.parent !is Subroutine)
-            errors.err("currently defer is only supported in subroutine scope, not in nested scopes", defer.position)
     }
 
     override fun visit(expr: BinaryExpression) {

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

@@ -60,15 +60,16 @@ private fun setDeferMasks(program: PtProgram, errors: IErrorReporter): Map<PtSub
 
         for((deferIndex, defer) in defers.withIndex()) {
             // replace the defer statement with one that enables the bit in the mask for this defer
-            val idx = defer.parent.children.indexOf(defer)
+            val scope = defer.parent
+            val idx = scope.children.indexOf(defer)
             val enableDefer = PtAugmentedAssign("|=", defer.position)
             val target = PtAssignTarget(true, defer.position)
             target.add(PtIdentifier(sub.scopedName+"."+maskVarName, DataType.UBYTE, defer.position))
             enableDefer.add(target)
             // enable the bit for this defer (beginning with high bits so the handler can simply shift right to check them in reverse order)
             enableDefer.add(PtNumber(DataType.UBYTE, (1 shl (defers.size-1 - deferIndex)).toDouble(), defer.position))
-            enableDefer.parent = sub
-            sub.children[idx] = enableDefer
+            enableDefer.parent = scope
+            scope.children[idx] = enableDefer
         }
     }
 

compiler/test/ast/TestVariousCompilerAst.kt

@@ -850,6 +850,10 @@ main {
             cx16.r0++
             cx16.r1++
         }
+        
+        if cx16.r0==0 {
+            defer cx16.r1++
+        }
 
         if cx16.r0==0
             return cx16.r0+cx16.r1
@@ -863,7 +867,24 @@ main {
 
         // check the desugaring of the defer statements
         (sub.children[0] as PtVariable).name shouldBe "p8v_prog8_defers_mask"
-        val ifelse = sub.children[3] as PtIfElse
+
+        val firstDefer = sub.children[2] as PtAugmentedAssign
+        firstDefer.operator shouldBe "|="
+        firstDefer.target.identifier?.name shouldBe "p8b_main.p8s_test.p8v_prog8_defers_mask"
+        firstDefer.value.asConstInteger() shouldBe 4
+
+        val firstIf = sub.children[3] as PtIfElse
+        val deferInIf = firstIf.ifScope.children[0] as PtAugmentedAssign
+        deferInIf.operator shouldBe "|="
+        deferInIf.target.identifier?.name shouldBe "p8b_main.p8s_test.p8v_prog8_defers_mask"
+        deferInIf.value.asConstInteger() shouldBe 2
+
+        val lastDefer = sub.children[5] as PtAugmentedAssign
+        lastDefer.operator shouldBe "|="
+        lastDefer.target.identifier?.name shouldBe "p8b_main.p8s_test.p8v_prog8_defers_mask"
+        lastDefer.value.asConstInteger() shouldBe 1
+
+        val ifelse = sub.children[4] as PtIfElse
         val ifscope = ifelse.ifScope.children[0] as PtNodeGroup
         val ifscope_push = ifscope.children[0] as PtFunctionCall
         val ifscope_defer = ifscope.children[1] as PtFunctionCall
@@ -871,10 +892,11 @@ main {
         ifscope_defer.name shouldBe "p8b_main.p8s_test.p8s_prog8_invoke_defers"
         ifscope_push.name shouldBe "sys.pushw"
         (ifscope_return.value as PtFunctionCall).name shouldBe "sys.popw"
-        val ending = sub.children[5] as PtFunctionCall
+
+        val ending = sub.children[6] as PtFunctionCall
         ending.name shouldBe "p8b_main.p8s_test.p8s_prog8_invoke_defers"
-        sub.children[6] shouldBe instanceOf<PtReturn>()
-        val handler = sub.children[7] as PtSub
+        sub.children[7] shouldBe instanceOf<PtReturn>()
+        val handler = sub.children[8] as PtSub
         handler.name shouldBe "p8s_prog8_invoke_defers"
     }
 }

docs/source/programming.rst

@@ -891,9 +891,9 @@ and the logic is declared very close to the spot where the allocation of the res
 It's possible to write a defer for a block of statements, but the advice is to keep such cleanup code as simple and short as possible.
 
 .. caution::
-    Defers only work for subroutines that are written as regular Prog8 code.
-    If a piece of inlined assembly somehow causes the routine to exit, the compiler cannot detect this.
-    Defers will not be handled in such cases.
+    Defers only work for subroutines that are written in regular Prog8 code.
+    If a piece of inlined assembly somehow causes the routine to exit, the compiler cannot detect this,
+    and defers won't be handled in such cases.
 
 
 Library routines and builtin functions

examples/test.p8

@@ -1,13 +1,28 @@
-%import floats
+%import textio
 %option no_sysinit
 %zeropage basicsafe
 
 main {
-
-    extsub $8000 = routine(float xx @FAC1, float yy @FAC2)
-
     sub start() {
-        @($8000) = $60
-        routine(1.234, 2.3445)
+        defer initial()
+        if allocate() {
+            txt.print("1")
+            defer deallocate()
+            txt.print("2")
+            return
+        }
+        txt.print("3")
+    }
+
+    sub allocate() -> bool {
+        txt.print("allocate\n")
+        return true
+    }
+
+    sub initial() {
+        txt.print("initial\n")
+    }
+    sub deallocate() {
+        txt.print("deallocate\n")
     }
 }