From 6a9acf0073326fc18ca4342b6c2c715e64b6918c Mon Sep 17 00:00:00 2001
From: yann300 <yann.levreau@gmail.com>
Date: Tue, 13 Feb 2024 13:29:39 +0100
Subject: [PATCH] fix uml gen

---
 apps/remix-ide/src/app/plugins/solidity-umlgen.tsx | 7 ++++---
 libs/remix-ui/solidity-uml-gen/src/types/index.ts  | 3 ++-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/apps/remix-ide/src/app/plugins/solidity-umlgen.tsx b/apps/remix-ide/src/app/plugins/solidity-umlgen.tsx
index ebbb7b9bb2..3fd6b6bd51 100644
--- a/apps/remix-ide/src/app/plugins/solidity-umlgen.tsx
+++ b/apps/remix-ide/src/app/plugins/solidity-umlgen.tsx
@@ -11,6 +11,7 @@ import vizRenderStringSync from '@aduh95/viz.js/sync'
 import {PluginViewWrapper} from '@remix-ui/helper'
 import {customAction} from '@remixproject/plugin-api'
 import {ClassOptions} from 'sol2uml/lib/converterClass2Dot'
+import type {CompilerInput} from '@remix-project/remix-solidity'
 const parser = (window as any).SolidityParser
 
 const _paq = (window._paq = window._paq || [])
@@ -74,7 +75,7 @@ export class SolidityUmlGen extends ViewPlugin implements ISolidityUmlGen {
       try {
         if (data.sources && Object.keys(data.sources).length > 1) {
           // we should flatten first as there are multiple asts
-          result = await this.flattenContract(source, file, data)
+          result = await this.flattenContract(source, file, data, JSON.parse(input))
         }
         const ast = result.length > 1 ? parser.parse(result) : parser.parse(source.sources[file].content)
         this.umlClasses = convertAST2UmlClasses(ast, this.currentFile)
@@ -142,8 +143,8 @@ export class SolidityUmlGen extends ViewPlugin implements ISolidityUmlGen {
    * and assigns to a local property
    * @returns {Promise<string>}
    */
-  async flattenContract(source: any, filePath: string, data: any) {
-    const result = await this.call('contractflattener', 'flattenContract', source, filePath, data)
+  async flattenContract(source: any, filePath: string, data: any, input: CompilerInput) {
+    const result = await this.call('contractflattener', 'flattenContract', source, filePath, data, input)
     return result
   }
 
diff --git a/libs/remix-ui/solidity-uml-gen/src/types/index.ts b/libs/remix-ui/solidity-uml-gen/src/types/index.ts
index 4359302642..a90b9340d6 100644
--- a/libs/remix-ui/solidity-uml-gen/src/types/index.ts
+++ b/libs/remix-ui/solidity-uml-gen/src/types/index.ts
@@ -1,5 +1,6 @@
 import { ViewPlugin } from '@remixproject/engine-web'
 import { customAction } from '@remixproject/plugin-api'
+import type {CompilerInput} from '@remix-project/remix-solidity'
 import React from 'react'
 
 export interface ISolidityUmlGen extends ViewPlugin {
@@ -17,7 +18,7 @@ export interface ISolidityUmlGen extends ViewPlugin {
     updateComponent(state: any): JSX.Element
     setDispatch(dispatch: React.Dispatch<any>): void
     generateCustomAction(action: customAction): Promise<void>
-    flattenContract (source: any, filePath: string, data: any): Promise<string>
+    flattenContract (source: any, filePath: string, data: any, input: CompilerInput): Promise<string>
     hideSpinner(): void
     renderComponent (): void
     triggerGenerateUml: boolean