From 5c6815f658d8eb54507be732c53b86c030db6425 Mon Sep 17 00:00:00 2001
From: Eric Lau <ericglau@outlook.com>
Date: Wed, 29 Nov 2023 10:06:52 -0500
Subject: [PATCH] Sort imports

---
 packages/core-cairo/src/print.ts | 33 ++++++++++++++++++++++++--------
 1 file changed, 25 insertions(+), 8 deletions(-)

diff --git a/packages/core-cairo/src/print.ts b/packages/core-cairo/src/print.ts
index e688fdbc0..c3171514b 100644
--- a/packages/core-cairo/src/print.ts
+++ b/packages/core-cairo/src/print.ts
@@ -75,17 +75,34 @@ function withSemicolons(lines: string[]): string[] {
 }
 
 function printImports(contract: Contract) {
-  const lines = [];
-  for (const component of contract.components) {
-    lines.push(`use ${component.path}::${component.name}`);
-  }
-  for (const standaloneImport of contract.standaloneImports) {
-    lines.push(`use ${standaloneImport}`);
+  const { ozImports, otherImports, superImports } = sortImports(contract);
+
+  const lines: string[] = [];
+  ozImports.forEach(i => lines.push(`use ${i}`));
+  otherImports.forEach(i => lines.push(`use ${i}`));
+  superImports.forEach(i => lines.push(`use ${i}`));
+  return withSemicolons(lines);
+}
+
+function sortImports(contract: Contract) {
+  const componentImports = contract.components.flatMap(c => `${c.path}::${c.name}`);
+  const combined = componentImports.concat(contract.standaloneImports);
+
+  const ozImports = [];
+  const otherImports = [];
+  const superImports = [];
+
+  for (const importStatement of combined) {
+    if (importStatement.startsWith('openzeppelin')) {
+      ozImports.push(importStatement);
+    } else {
+      otherImports.push(importStatement);
+    }
   }
   if (contract.superVariables.length > 0) {
-    lines.push(`use super::{${contract.superVariables.map(v => v.name).join(', ')}}`);
+    superImports.push(`super::{${contract.superVariables.map(v => v.name).join(', ')}}`);
   }
-  return withSemicolons(lines);
+  return { ozImports, otherImports, superImports };
 }
 
 function printComponentDeclarations(contract: Contract) {