From 9f78985e4567ff2c5d300c6bc21353d8d7f94501 Mon Sep 17 00:00:00 2001
From: Abhishek Aggarwal <aaggarwal@mozilla.com>
Date: Mon, 10 May 2021 19:09:18 +0200
Subject: [PATCH] JS bindings for vocabularies as bytes

---
 wasm/bindings/TranslationModelBindings.cpp | 28 ++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/wasm/bindings/TranslationModelBindings.cpp b/wasm/bindings/TranslationModelBindings.cpp
index 41b9c2e..3fd2c8b 100644
--- a/wasm/bindings/TranslationModelBindings.cpp
+++ b/wasm/bindings/TranslationModelBindings.cpp
@@ -24,12 +24,36 @@ EMSCRIPTEN_BINDINGS(aligned_memory) {
     .function("size", &marian::bergamot::AlignedMemory::size)
 	  .function("getByteArrayView", &getByteArrayView)
     ;
+
+    register_vector<marian::bergamot::AlignedMemory*>("AlignedMemoryList");
+}
+
+std::vector<std::shared_ptr<marian::bergamot::AlignedMemory>>
+prepareVocabsSmartMemories(std::vector<marian::bergamot::AlignedMemory*>& vocabsMemories) {
+  auto sourceVocabMemory = std::make_shared<marian::bergamot::AlignedMemory>(std::move(*(vocabsMemories[0])));
+  std::vector<std::shared_ptr<marian::bergamot::AlignedMemory>> vocabsSmartMemories;
+  vocabsSmartMemories.push_back(sourceVocabMemory);
+  // When source and target vocab files are same, only one memory object is passed in vocabsMemories
+  // to avoid double memory allocation for the same file. However, the constructor of the TranslationModel
+  // class still expects 2 entries where each entry has the shared ownership of a single AlignedMemory object.
+  if (vocabsMemories.size() == 2) {
+    auto targetVocabMemory = std::make_shared<marian::bergamot::AlignedMemory>(std::move(*(vocabsMemories[1])));
+    vocabsSmartMemories.push_back(std::move(targetVocabMemory));
+  }
+  else {
+    vocabsSmartMemories.push_back(sourceVocabMemory);
+  }
+  return vocabsSmartMemories;
 }
 
 TranslationModel* TranslationModelFactory(const std::string &config,
                                           marian::bergamot::AlignedMemory* modelMemory,
-                                          marian::bergamot::AlignedMemory* shortlistMemory) {
-  return new TranslationModel(config, std::move(*modelMemory), std::move(*shortlistMemory));
+                                          marian::bergamot::AlignedMemory* shortlistMemory,
+                                          std::vector<marian::bergamot::AlignedMemory*> uniqueVocabsMemories) {
+  return new TranslationModel(config,
+                              std::move(*modelMemory),
+                              std::move(*shortlistMemory),
+                              std::move(prepareVocabsSmartMemories(uniqueVocabsMemories)));
 }
 
 EMSCRIPTEN_BINDINGS(translation_model) {