Fixed MessageCollector to handle async return messages correctly.

src/parser/MessageCollector.ts

@@ -13,8 +13,13 @@ export class MessageCollector extends sequenceParserListener {
     this._addOwnedMessage(OwnableMessageType.AsyncMessage)(ctx);
   enterCreation = (ctx: any) =>
     this._addOwnedMessage(OwnableMessageType.CreationMessage)(ctx);
-  enterRet = (ctx: any) =>
+  enterRet = (ctx: any) => {
+    if (ctx.asyncMessage()) {
+      // it will visit the asyncMessage later
+      return;
+    }
     this._addOwnedMessage(OwnableMessageType.ReturnMessage)(ctx);
+  };
 
   private _addOwnedMessage = (type: OwnableMessageType) => (ctx: any) => {
     if (this.isBlind) {

src/positioning/MessageCollector.spec.ts

@@ -12,6 +12,18 @@ describe("MessageCollector", () => {
   };
 
   describe("Return", () => {
+    it("@return A->B:m", () => {
+      const code = `@return A->B:m`;
+
+      expect(collectMessages(code)).toStrictEqual([
+        {
+          from: "A",
+          signature: "m",
+          to: "B",
+          type: OwnableMessageType.AsyncMessage,
+        },
+      ]);
+    });
     it("should handle one return message", () => {
       const code = `
         return result

test/unit/parser/participant/Participants.Order.spec.ts

@@ -34,13 +34,20 @@ import {
   _STARTER_,
   OrderedParticipants,
 } from "../../../../src/parser/OrderedParticipants";
+import { expect } from "vitest";
 
 function getFlattenedParticipants(code: string) {
   const rootContext = RootContext(code);
   return OrderedParticipants(rootContext);
 }
 
 describe("Participants.Order", () => {
+  it("@return", () => {
+    expect(getFlattenedParticipants("@return A->B:m")).toEqual([
+      { name: "A", left: "" },
+      { name: "B", left: "A" },
+    ]);
+  });
   it("should return the order of participants", () => {
     expect(getFlattenedParticipants("A as A1 B C.m")).toEqual([
       { name: _STARTER_, left: "" },

test/unit/parser/to-collector.spec.js

@@ -4,6 +4,12 @@ import ToCollector from "../../../src/parser/ToCollector";
 import { expect } from "vitest";
 import { blankParticipant } from "../../../src/parser/Participants";
 
+test("@return", () => {
+  let participants = getParticipants("@return A->B.m");
+  expect(participants.Size()).toBe(2);
+  expect(participants.First().name).toBe("A");
+});
+
 test("if block", () => {
   let participants = getParticipants("if(x) { A->B.m }");
   expect(participants.Size()).toBe(2);