From d6919a02b1d03903bc0ed75896d34bce44b7ea80 Mon Sep 17 00:00:00 2001
From: Chandler May <457238+ccmaymay@users.noreply.github.com>
Date: Mon, 16 Oct 2023 20:40:45 -0400
Subject: [PATCH] Return added Communication-level references for TS users.
 (#13)

---
 src_nodejs/communication_fu.d.ts     | 14 ++++++--
 src_nodejs/communication_fu.js       | 24 +++++++++++---
 test_nodejs/communication_fu_test.js | 48 ++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 8 deletions(-)
 create mode 100644 test_nodejs/communication_fu_test.js

diff --git a/src_nodejs/communication_fu.d.ts b/src_nodejs/communication_fu.d.ts
index 6ae71ee..bc1a474 100644
--- a/src_nodejs/communication_fu.d.ts
+++ b/src_nodejs/communication_fu.d.ts
@@ -1,12 +1,20 @@
 import './communication_types';
 import {UUID} from './uuid_types';
-import {Entity, EntityMentionSet} from './entities_types';
+import {Entity, EntityMention, EntityMentionSet} from './entities_types';
 import {Section, Sentence, Tokenization, Token} from './structure_types';
-import {SituationMention} from './situations_types';
+import {Situation, SituationMention} from './situations_types';
 
 declare module "./communication_types" {
   interface Communication {
-    addInternalReferences(): void;
+    addInternalReferences(): {
+      entityForUUID: {[index: string]: Entity},
+      entityMentionForUUID: {[index: string]: EntityMention},
+      sectionForUUID: {[index: string]: Section},
+      sentenceForUUID: {[index: string]: Sentence},
+      situationForUUID: {[index: string]: Situation},
+      situationMentionForUUID: {[index: string]: SituationMention},
+      tokenizationForUUID: {[index: string]: Tokenization},
+    };
     getEntityForEntityMentionUUID(uuid: UUID): Entity?;
     getEntityMentionSetWithToolname(toolname: string): EntityMentionSet;
     getEntityMentionWithUUID(uuid: UUID): EntityMention?;
diff --git a/src_nodejs/communication_fu.js b/src_nodejs/communication_fu.js
index e59e098..d71b7e4 100644
--- a/src_nodejs/communication_fu.js
+++ b/src_nodejs/communication_fu.js
@@ -1,5 +1,5 @@
 const {Communication} = require('./communication_types');
-const {Thrift} = require('thrift');
+const {TBufferedTransport, TJSONProtocol} = require('thrift');
 
 /**
  * Adds internal references between data structures contained in Communication.
@@ -42,6 +42,10 @@ const {Thrift} = require('thrift');
  *     - tokenization: reference to the Tokenization corresponding to tokenizationId
  *
  * @function concrete.communication.Communication.prototype.addInternalReferences
+ * @returns {object} Object with references to Communication-level maps (entityForUUID, entityMentionForUUID,
+ *                   sectionForUUID, sentenceForUUID, situationForUUID, situationMentionForUUID, tokenizationForUUID),
+ *                   mainly intended for TypeScript consumers.  JavaScript consumers can access these maps directly
+ *                   on the Communication.
  */
 Communication.prototype.addInternalReferences = function() {
   this.entityForUUID = {};
@@ -144,6 +148,16 @@ Communication.prototype.addInternalReferences = function() {
       situation.situationSet = situationSet;
     });
   });
+
+  return {
+    entityForUUID: this.entityForUUID,
+    entityMentionForUUID: this.entityMentionForUUID,
+    sectionForUUID: this.sectionForUUID,
+    sentenceForUUID: this.sentenceForUUID,
+    situationForUUID: this.situationForUUID,
+    situationMentionForUUID: this.situationMentionForUUID,
+    tokenizationForUUID: this.tokenizationForUUID,
+  };
 };
 
 
@@ -488,8 +502,8 @@ Communication.prototype.initFromTJSONProtocolObject = function(commJSONObject) {
  */
 Communication.prototype.initFromTJSONProtocolString = function(commJSONString) {
   var commJSONObject = JSON.parse(commJSONString);
-  var transport = new Thrift.Transport();
-  var protocol = new Thrift.TJSONProtocol(transport);
+  var transport = new TBufferedTransport();
+  var protocol = new TJSONProtocol(transport);
 
   // The values for these protocol object fields was determined by
   // mucking around with the JavaScript debugger to figure out how
@@ -524,8 +538,8 @@ Communication.prototype.toTJSONProtocolObject = function() {
  * @returns {string}
  */
 Communication.prototype.toTJSONProtocolString = function() {
-  var transport = new Thrift.Transport();
-  var protocol = new Thrift.TJSONProtocol(transport);
+  var transport = new TBufferedTransport();
+  var protocol = new TJSONProtocol(transport);
   protocol.tpos = [];
   protocol.tstack = [];
   this.write(protocol);
diff --git a/test_nodejs/communication_fu_test.js b/test_nodejs/communication_fu_test.js
new file mode 100644
index 0000000..3d1ad8f
--- /dev/null
+++ b/test_nodejs/communication_fu_test.js
@@ -0,0 +1,48 @@
+const expect = require("chai").expect;
+
+const concrete = require("../dist_nodejs");
+const commJSONData = require('../test/fixtures/dog-bites-man.concrete.json');
+
+describe("addInternalReferences", function() {
+  it("Returns Communication-level maps", function() {
+    const comm = new concrete.communication.Communication();
+    comm.initFromTJSONProtocolObject(commJSONData);
+    const refs = comm.addInternalReferences();
+    expect(refs).to.have.all.keys([
+      "entityForUUID",
+      "entityMentionForUUID",
+      "sectionForUUID",
+      "sentenceForUUID",
+      "situationForUUID",
+      "situationMentionForUUID",
+      "tokenizationForUUID"
+    ]);
+    comm.sectionList.forEach((section) => {
+      expect(refs.sectionForUUID[section.uuid.uuidString]).to.equal(section);
+      section.sentenceList.forEach((sentence) => {
+        expect(refs.sentenceForUUID[sentence.uuid.uuidString]).to.equal(sentence);
+        expect(refs.tokenizationForUUID[sentence.tokenization.uuid.uuidString]).to.equal(sentence.tokenization);
+      });
+    });
+    comm.entitySetList.forEach((entitySet) =>
+      entitySet.entityList.forEach((entity) =>
+        expect(refs.entityForUUID[entity.uuid.uuidString]).to.equal(entity)
+      )
+    );
+    comm.entityMentionSetList.forEach((entityMentionSet) =>
+      entityMentionSet.mentionList.forEach((mention) =>
+        expect(refs.entityMentionForUUID[mention.uuid.uuidString]).to.equal(mention)
+      )
+    );
+    comm.situationSetList.forEach((situationSet) =>
+      situationSet.situationList.forEach((situation) =>
+        expect(refs.situationForUUID[situation.uuid.uuidString]).to.equal(situation)
+      )
+    );
+    comm.situationMentionSetList.forEach((situationMentionSet) =>
+      situationMentionSet.mentionList.forEach((mention) =>
+        expect(refs.situationMentionForUUID[mention.uuid.uuidString]).to.equal(mention)
+      )
+    );
+  });
+});