From 7a645fc4b443b0bbdf2805e09cce19c93f63abc6 Mon Sep 17 00:00:00 2001
From: root <melchert@stanford.edu>
Date: Fri, 24 May 2024 10:52:06 -0700
Subject: [PATCH] Returning dict for better input and output unrolling handling

---
 src/CodeGen_Pono.cpp | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/CodeGen_Pono.cpp b/src/CodeGen_Pono.cpp
index 66e73ffbe..fbf49729f 100644
--- a/src/CodeGen_Pono.cpp
+++ b/src/CodeGen_Pono.cpp
@@ -362,17 +362,17 @@ string associated_provide_name(Stmt s, string call_name) {
                 testbench_stream << ")\n";
 
                 do_indent();
-                testbench_stream << "return ";
+                testbench_stream << "return {";
                 for (size_t i = 0; i < args.size(); i++) {
                     if (args[i].is_stencil) {
-                        testbench_stream << print_name(args[i].name);
+                        testbench_stream << "\"" << print_name(args[i].name) << "\": " << print_name(args[i].name);
 
                         if (i < args.size() - 1) {
                             testbench_stream << ", ";
                         }
                     }
                 }
-                testbench_stream << "\n";
+                testbench_stream << "}\n";
                 
 
 
@@ -387,16 +387,22 @@ string associated_provide_name(Stmt s, string call_name) {
                 testbench_stream << "solver = Solver()\n";
                 do_indent();
                 
+                testbench_stream << "inputs_and_outputs_dict = create_app(solver)\n";
+
+                // do_indent();
+
                 for (size_t i = 0; i < args.size(); i++) {
                     if (args[i].is_stencil) {
-                        testbench_stream << print_name(args[i].name);
+                        do_indent();
+                        testbench_stream << print_name(args[i].name) << " = inputs_and_outputs_dict[\"" << print_name(args[i].name) << "\"]\n";
 
-                        if (i < args.size() - 1) {
-                            testbench_stream << ", ";
-                        }
+                        // if (i < args.size() - 1) {
+                        //     testbench_stream << "\n";
+                        //     do_indent();
+                        // }
                     }
                 }
-                testbench_stream << " = create_app(solver)\n";
+                
 
                 // read input image(s)
                 for (size_t i = 0; i < args.size() - 1; i++) {