feat(cfg): mid marker for for loops (#1610)

feat(cfg): mid marker for for loops

Author: EagleoutIce

src/control-flow/extract-cfg.ts

@@ -221,7 +221,7 @@ function cfgWhile(whileLoop: RWhileLoop<ParentInformation>, condition: ControlFl
 
 function cfgFor(forLoop: RForLoop<ParentInformation>, variable: ControlFlowInformation, vector: ControlFlowInformation, body: ControlFlowInformation): ControlFlowInformation {
 	const graph = variable.graph;
-	graph.addVertex({ id: forLoop.info.id, type: identifyMayStatementType(forLoop), exit: [forLoop.info.id + '-exit'] });
+	graph.addVertex({ id: forLoop.info.id, type: identifyMayStatementType(forLoop), exit: [forLoop.info.id + '-exit'], mid: [forLoop.info.id + '-head'] });
 
 	graph.merge(vector.graph);
 	graph.merge(body.graph);
@@ -236,10 +236,13 @@ function cfgFor(forLoop: RForLoop<ParentInformation>, variable: ControlFlowInfor
 		}
 	}
 
+	graph.addVertex({ id: forLoop.info.id + '-head', type: CfgVertexType.MidMarker, root: forLoop.info.id, kind: 'head' });
+
 	for(const exit of variable.exitPoints) {
-		for(const entry of body.entryPoints) {
-			graph.addEdge(entry, exit, { label: CfgEdgeType.Cd, when: RTrue, caused: forLoop.info.id });
-		}
+		graph.addEdge(forLoop.info.id + '-head', exit, { label: CfgEdgeType.Fd });
+	}
+	for(const entry of body.entryPoints) {
+		graph.addEdge(entry, forLoop.info.id + '-head', { label: CfgEdgeType.Cd, when: RTrue, caused: forLoop.info.id });
 	}
 
 	for(const next of [...body.nexts, ...body.exitPoints]) {

test/functionality/control-flow/control-flow-graph.test.ts

@@ -216,6 +216,15 @@ describe('Control Flow Graph', withTreeSitter(parser => {
 					.addEdge('10-exit', 9, { label: CfgEdgeType.Fd })
 					.addEdge(11, '10-exit', { label: CfgEdgeType.Fd })
 			});
+
+			assertCfg(parser, 'for (i in v) b', {
+				entryPoints: [ '5' ],
+				exitPoints:  [ '5-exit' ],
+				graph:       new ControlFlowGraph()
+					.addVertex({ id: '4-head', kind: 'head', type: CfgVertexType.MidMarker, root: 4 })
+			}, {
+				expectIsSubgraph: true
+			});
 		});
 
 		describe('function calls', () => {