--- a/InstrumentingCompiler.st	Mon May 03 19:41:23 2010 +0200
+++ b/InstrumentingCompiler.st	Mon May 03 19:41:25 2010 +0200
@@ -14,8 +14,8 @@
 	privateIn:InstrumentingCompiler
 !
 
-InstrumentingCompiler::InstrumentationInfo subclass:#BlockExecutionInfo
-	instanceVariableNames:'startPosition endPosition count'
+InstrumentingCompiler::InstrumentationInfo subclass:#StatementExecutionInfo
+	instanceVariableNames:'startPosition count'
 	classVariableNames:''
 	poolDictionaries:''
 	privateIn:InstrumentingCompiler
@@ -42,6 +42,13 @@
 	privateIn:InstrumentingCompiler::MethodInvocationInfo::MethodInvocationInfoPerReceiverClass
 !
 
+InstrumentingCompiler::StatementExecutionInfo subclass:#BlockExecutionInfo
+	instanceVariableNames:'endPosition'
+	classVariableNames:''
+	poolDictionaries:''
+	privateIn:InstrumentingCompiler
+!
+
 !InstrumentingCompiler class methodsFor:'documentation'!
 
 documentation
@@ -190,6 +197,27 @@
     ^ InstrumentedMethod
 
     "Created: / 27-04-2010 / 12:17:22 / cg"
+!
+
+statementCounterBefore:aStatementNode
+    |countCode statementEntryInfo|
+
+    blockInvocationInfo isNil ifTrue:[
+        blockInvocationInfo := OrderedCollection new.
+    ].
+    statementEntryInfo := StatementExecutionInfo new cleanInfo.
+    statementEntryInfo startPosition:aStatementNode startPosition.
+    blockInvocationInfo add:statementEntryInfo.
+
+    countCode := 
+        StatementNode 
+            expression:(MessageNode 
+                            receiver:(ConstantNode value:statementEntryInfo)
+                            selector:#entry:
+                            arg:(VariableNode type:#ThisContext context:contextToEvaluateIn)).
+
+    countCode nextStatement:aStatementNode.
+    ^ countCode
 ! !
 
 !InstrumentingCompiler methodsFor:'code generation-hooks'!
@@ -236,6 +264,39 @@
             codeForSideEffectOn:codeStream inBlock:nil for:self.
 
     "Modified: / 27-04-2010 / 11:50:48 / cg"
+!
+
+statementListRewriteHookFor:firstStatement
+    "/ care for blocks which return...
+    |thisStatement prevStatement countStatement mustInsertExecutionInfo|
+
+    mustInsertExecutionInfo := false.
+    prevStatement := nil.
+
+    thisStatement := firstStatement.
+    thisStatement isNil ifTrue:[^ firstStatement].
+
+    [true] whileTrue:[
+        mustInsertExecutionInfo ifTrue:[
+            "/ prev-stat had a return in it (a block with a return);
+            "/ insert a statement-entry here.
+            thisStatement notNil ifTrue:[
+                countStatement := self statementCounterBefore:thisStatement.
+                prevStatement nextStatement:countStatement.
+                mustInsertExecutionInfo := false.
+                prevStatement := countStatement.
+            ]
+        ].
+        thisStatement isNil ifTrue:[^ firstStatement].
+
+        thisStatement isReturnNode ifFalse:[
+            thisStatement containsReturn ifTrue:[
+                mustInsertExecutionInfo := true    
+            ]
+        ].
+        prevStatement := thisStatement.
+        thisStatement := thisStatement nextStatement.
+    ].
 ! !
 
 !InstrumentingCompiler::InstrumentationInfo class methodsFor:'cleanup'!
@@ -279,7 +340,7 @@
     "Created: / 27-04-2010 / 12:06:48 / cg"
 ! !
 
-!InstrumentingCompiler::BlockExecutionInfo methodsFor:'accessing'!
+!InstrumentingCompiler::StatementExecutionInfo methodsFor:'accessing'!
 
 callCount
     ^ count
@@ -307,32 +368,23 @@
     "Created: / 23-06-2006 / 13:31:28 / cg"
 !
 
-endPosition
-    ^ endPosition
-
-    "Created: / 28-04-2010 / 15:57:14 / cg"
-!
-
 startPosition
     ^ startPosition
 
     "Created: / 28-04-2010 / 15:54:26 / cg"
 !
 
-startPosition:startArg endPosition:endArg
-    startPosition := startArg.
-    endPosition := endArg.
-
-    "Created: / 28-04-2010 / 15:54:47 / cg"
+startPosition:something
+    startPosition := something.
 ! !
 
-!InstrumentingCompiler::BlockExecutionInfo methodsFor:'cleanup'!
+!InstrumentingCompiler::StatementExecutionInfo methodsFor:'cleanup'!
 
 cleanInfo
     count := 0
 ! !
 
-!InstrumentingCompiler::BlockExecutionInfo methodsFor:'instrumentation calls'!
+!InstrumentingCompiler::StatementExecutionInfo methodsFor:'instrumentation calls'!
 
 entry:callingContext
     "probe entry from instrumented code;
@@ -348,7 +400,7 @@
     "Modified: / 27-04-2010 / 14:03:29 / cg"
 ! !
 
-!InstrumentingCompiler::BlockExecutionInfo methodsFor:'queries'!
+!InstrumentingCompiler::StatementExecutionInfo methodsFor:'queries'!
 
 hasBeenExecuted
     ^ count > 0
@@ -574,8 +626,23 @@
     "Created: / 27-04-2010 / 18:17:27 / cg"
 ! !
 
+!InstrumentingCompiler::BlockExecutionInfo methodsFor:'accessing'!
+
+endPosition
+    ^ endPosition
+
+    "Created: / 28-04-2010 / 15:57:14 / cg"
+!
+
+startPosition:startArg endPosition:endArg
+    startPosition := startArg.
+    endPosition := endArg.
+
+    "Created: / 28-04-2010 / 15:54:47 / cg"
+! !
+
 !InstrumentingCompiler class methodsFor:'documentation'!
 
 version_CVS
-    ^ '$Header: /cvs/stx/stx/libcomp/InstrumentingCompiler.st,v 1.4 2010-04-30 09:59:44 cg Exp $'
+    ^ '$Header: /cvs/stx/stx/libcomp/InstrumentingCompiler.st,v 1.5 2010-05-03 17:41:25 cg Exp $'
 ! !