sol2uml 2.1.0 → 2.1.3

package/README.md

@@ -2,17 +2,19 @@
 
 [![npm version](https://badge.fury.io/js/sol2uml.svg)](https://badge.fury.io/js/sol2uml)
 
-[Unified Modeling Language (UML)](https://en.wikipedia.org/wiki/Unified_Modeling_Language) [class diagram](https://en.wikipedia.org/wiki/Class_diagram) generator for [Solidity](https://solidity.readthedocs.io/) contracts.
+A visualisation tool for [Solidity](https://solidity.readthedocs.io/) contracts featuring:
+1. [Unified Modeling Language (UML)](https://en.wikipedia.org/wiki/Unified_Modeling_Language) [class diagram](https://en.wikipedia.org/wiki/Class_diagram) generator for Solidity contracts.
+2. Contract storage layout diagrams.
 
-Open Zeppelin's ERC20 token contracts generated from [version 2.5.1](https://github.com/OpenZeppelin/openzeppelin-solidity/tree/v2.5.1/contracts/token/ERC20)
+UML class diagram of Open Zeppelin's ERC20 token contracts generated from [version 2.5.1](https://github.com/OpenZeppelin/openzeppelin-solidity/tree/v2.5.1/contracts/token/ERC20)
 ![Open Zeppelin ERC20](./examples/OpenZeppelinERC20.svg)
 
 See more contract diagrams [here](./examples/README.md).
 
-USDC storage slots from the [verified source code](https://etherscan.io/address/0xa2327a938febf5fec13bacfb16ae10ecbc4cbdcf#code) on Etherscan.
-![USDC](./examples/storage/usdc.png)
+Storage layout diagram of USDC's [verified source code](https://etherscan.io/address/0xa2327a938febf5fec13bacfb16ae10ecbc4cbdcf#code) on Etherscan.
+![USDC](./examples/storage/usdcData.svg)
 
-See more storage slot diagrams [here](./examples/storage/README.md).
+See more contract storage diagram examples [here](./examples/storage/README.md).
 
 # Install
 
@@ -42,16 +44,16 @@ npm ls sol2uml -g
 ## Command Line Interface (CLI)
 
 ```
-$ sol2uml --help
 Usage: sol2uml [subcommand] <options>
 The three subcommands:
 * class:    Generates a UML class diagram from Solidity source code. default
 * storage:  Generates a diagram of a contract's storage slots.
-* flatten:  Pulls verified source files from a Blockchain explorer into one, flat, local Solidity file.
+* flatten:  Merges verified source files from a Blockchain explorer into one local file.
 
 The Solidity code can be pulled from verified source code on Blockchain explorers like Etherscan or from local Solidity files.
 
 Options:
+  -V, --version                                output the version number
   -sf, --subfolders <value>                    number of subfolders that will be recursively searched for Solidity files. (default: all)
   -f, --outputFormat <value>                   output file format. (choices: "svg", "png", "dot", "all", default: "svg")
   -o, --outputFileName <value>                 output file name
@@ -71,7 +73,6 @@ Commands:
 ### Class usage
 
 ```
-$sol2uml class --help
 Usage: sol2uml class <fileFolderAddress> [options]
 
 Generates UML diagrams from Solidity source code.
@@ -108,20 +109,19 @@ Options:
 ### Storage usage
 
 ```
-Usage: sol2uml storage [options] <fileFolderAddress>
-
-Visually display a contract's storage slots.
+Usage: sol2uml storage <fileFolderAddress> [options]
 
 WARNING: sol2uml does not use the Solidity compiler so may differ with solc. A known example is fixed-sized arrays declared with an expression will fail to be sized.
 
+Visually display a contract's storage slots.
+
 Arguments:
   fileFolderAddress        file name, base folder or contract address
 
 Options:
   -c, --contract <name>    Contract name in local Solidity files. Not needed when using an address as the first argument as the contract name can be derived from Etherscan.
   -d, --data               Gets the values in the storage slots from an Ethereum node. (default: false)
-  -s, --storage <address>  The address of the contract with the storage values. This will be different from the contract with the code if a proxy contract is used. This is not needed if `fileFolderAddress` is an address and
-                           the contract is not proxied.
+  -s, --storage <address>  The address of the contract with the storage values. This will be different from the contract with the code if a proxy contract is used. This is not needed if `fileFolderAddress` is an address and the contract is not proxied.
   -u, --url <url>          URL of the Ethereum node to get storage values if the `data` option is used. (default: "http://localhost:8545", env: NODE_URL)
   -bn, --block <number>    Block number to get the contract storage values from. (default: "latest")
   -h, --help               display help for command
@@ -130,17 +130,20 @@ Options:
 ### Flatten usage
 
 ```
-$sol2uml flatten --help
-Usage: sol2uml flatten [options] <contractAddress>
+Usage: sol2uml flatten <contractAddress> [options]
+
+In order for the merged code to compile, the following is done:
+1. File imports are commented out.
+2. "SPDX-License-Identifier" is renamed to "SPDX--License-Identifier".
+3. Contract dependencies are analysed so the files are merged in an order that will compile.
 
-get all verified source code for a contract from the Blockchain explorer into one local file
+Merges verified source files for a contract from a Blockchain explorer into one local file.
 
 Arguments:
-  contractAddress  Contract address
+  contractAddress  Contract address in hexadecimal format with a 0x prefix.
 
 Options:
   -h, --help       display help for command
-
 ```
 
 ## UML Class diagram examples

package/lib/associations.d.ts

@@ -1,2 +1,2 @@
 import { Association, UmlClass } from './umlClass';
-export declare const findAssociatedClass: (association: Association, sourceUmlClass: UmlClass, umlClasses: UmlClass[]) => UmlClass;
+export declare const findAssociatedClass: (association: Association, sourceUmlClass: UmlClass, umlClasses: UmlClass[]) => UmlClass | undefined;

package/lib/associations.js

@@ -1,31 +1,75 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.findAssociatedClass = void 0;
+const umlClass_1 = require("./umlClass");
 // Find the UML class linked to the association
 const findAssociatedClass = (association, sourceUmlClass, umlClasses) => {
-    return umlClasses.find((targetUmlClass) => {
-        return (
-        // class is in the same source file
+    let umlClass = umlClasses.find((targetUmlClass) => {
+        // is the source class link via the association to the target class?
+        if (isAssociated(association, sourceUmlClass, targetUmlClass))
+            return true;
+        // Not linked so now try linking to target under the node_modules folder.
+        // eg remove node_modules from node_modules/@openzeppelin/contracts-upgradeable/proxy/Initializable.sol
+        // is the target class under node_modules?
+        if (targetUmlClass.relativePath.match(/^node_modules\//)) {
+            // clone the target and updated absolutePath and relativePath so it's no longer under node_modules
+            const clonedTargetClass = new umlClass_1.UmlClass(targetUmlClass);
+            clonedTargetClass.absolutePath =
+                targetUmlClass.absolutePath.replace(/^node_modules\//, '');
+            clonedTargetClass.relativePath =
+                targetUmlClass.relativePath.replace(/^node_modules\//, '');
+            // is the source class link via the association to the target class?
+            return isAssociated(association, sourceUmlClass, clonedTargetClass);
+        }
+        // could not find a link from the source to target via the association
+        return false;
+    });
+    // If a link was found
+    if (umlClass)
+        return umlClass;
+    // Could not find a link so now need to recursively look at imports of imports
+    return findImplicitImport(association, sourceUmlClass, umlClasses);
+};
+exports.findAssociatedClass = findAssociatedClass;
+// Tests if source class can be linked to the target class via an association
+const isAssociated = (association, sourceUmlClass, targetUmlClass) => {
+    return (
+    // class is in the same source file
+    (association.targetUmlClassName === targetUmlClass.name &&
+        sourceUmlClass.absolutePath === targetUmlClass.absolutePath) ||
+        // imported classes with no explicit import names
         (association.targetUmlClassName === targetUmlClass.name &&
-            sourceUmlClass.absolutePath === targetUmlClass.absolutePath) ||
-            // imported classes with no explicit import names
-            (association.targetUmlClassName === targetUmlClass.name &&
-                sourceUmlClass.imports.find((i) => i.absolutePath === targetUmlClass.absolutePath &&
-                    i.classNames.length === 0)) ||
-            // imported classes with explicit import names or import aliases
             sourceUmlClass.imports.find((i) => i.absolutePath === targetUmlClass.absolutePath &&
-                i.classNames.find((importedClass) => 
-                // no import alias
+                i.classNames.length === 0)) ||
+        // imported classes with explicit import names or import aliases
+        sourceUmlClass.imports.find((i) => i.absolutePath === targetUmlClass.absolutePath &&
+            i.classNames.find((importedClass) => 
+            // no import alias
+            (association.targetUmlClassName ===
+                importedClass.className &&
+                importedClass.className === targetUmlClass.name &&
+                importedClass.alias == undefined) ||
+                // import alias
                 (association.targetUmlClassName ===
-                    importedClass.className &&
-                    importedClass.className ===
-                        targetUmlClass.name &&
-                    importedClass.alias == undefined) ||
-                    // import alias
-                    (association.targetUmlClassName ===
-                        importedClass.alias &&
-                        importedClass.className === targetUmlClass.name))));
-    });
+                    importedClass.alias &&
+                    importedClass.className === targetUmlClass.name))));
+};
+const findImplicitImport = (association, sourceUmlClass, umlClasses) => {
+    // Get all implicit imports. That is, imports that do not explicitly import contracts or interfaces.
+    const implicitImports = sourceUmlClass.imports.filter((i) => i.classNames.length === 0);
+    // For each implicit import
+    for (const importDetail of implicitImports) {
+        // Find a class with the same absolute path as the import so we can get the new imports
+        const newSourceUmlClass = umlClasses.find((c) => c.absolutePath === importDetail.absolutePath);
+        if (!newSourceUmlClass) {
+            // Could not find a class in the import file so just move onto the next loop
+            continue;
+        }
+        // TODO need to handle imports that use aliases as the association will not be found
+        const umlClass = (0, exports.findAssociatedClass)(association, newSourceUmlClass, umlClasses);
+        if (umlClass)
+            return umlClass;
+    }
+    return undefined;
 };
-exports.findAssociatedClass = findAssociatedClass;
 //# sourceMappingURL=associations.js.map

package/lib/converterAST2Classes.js

@@ -176,7 +176,7 @@ function parseContractDefinition(umlClass, node) {
                 });
                 // Is the variable a constant that could be used in declaring fixed sized arrays
                 if (variable.isDeclaredConst) {
-                    if (variable?.expression.type === 'NumberLiteral') {
+                    if (variable?.expression?.type === 'NumberLiteral') {
                         umlClass.constants.push({
                             name: variable.name,
                             value: parseInt(variable.expression.number),

package/lib/converterClasses2Storage.d.ts

@@ -1,7 +1,8 @@
 import { Attribute, UmlClass } from './umlClass';
 export declare enum StorageType {
-    Contract = 0,
-    Struct = 1
+    Contract = "Contract",
+    Struct = "Struct",
+    Array = "Array"
 }
 export interface Variable {
     id: number;
@@ -10,17 +11,22 @@ export interface Variable {
     byteSize: number;
     byteOffset: number;
     type: string;
-    variable: string;
+    dynamic: boolean;
+    variable?: string;
     contractName?: string;
-    values: string[];
-    structStorageId?: number;
+    noValue: boolean;
+    value?: string;
+    referenceStorageId?: number;
     enumId?: number;
 }
 export interface Storage {
     id: number;
     name: string;
     address?: string;
+    slotKey?: string;
     type: StorageType;
+    arrayLength?: number;
+    arrayDynamic?: boolean;
     variables: Variable[];
 }
 /**
@@ -31,6 +37,12 @@ export interface Storage {
  */
 export declare const addStorageValues: (url: string, contractAddress: string, storage: Storage, blockTag: string) => Promise<void>;
 export declare const convertClasses2Storages: (contractName: string, umlClasses: UmlClass[]) => Storage[];
-export declare const parseStructStorage: (attribute: Attribute, otherClasses: UmlClass[], storages: Storage[]) => Storage | undefined;
-export declare const calcStorageByteSize: (attribute: Attribute, umlClass: UmlClass, otherClasses: UmlClass[]) => number;
+export declare const parseReferenceStorage: (attribute: Attribute, umlClass: UmlClass, otherClasses: UmlClass[], storages: Storage[]) => Storage | undefined;
+export declare const calcStorageByteSize: (attribute: Attribute, umlClass: UmlClass, otherClasses: UmlClass[]) => {
+    size: number;
+    dynamic: boolean;
+};
 export declare const isElementary: (type: string) => boolean;
+export declare const calcSlotKey: (variable: Variable) => string | undefined;
+export declare const offsetStorageSlots: (storage: Storage, slots: number, storages: Storage[]) => void;
+export declare const findDimensionLength: (umlClass: UmlClass, dimension: string) => number;

package/lib/converterClasses2Storage.js

@@ -1,13 +1,16 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.isElementary = exports.calcStorageByteSize = exports.parseStructStorage = exports.convertClasses2Storages = exports.addStorageValues = exports.StorageType = void 0;
+exports.findDimensionLength = exports.offsetStorageSlots = exports.calcSlotKey = exports.isElementary = exports.calcStorageByteSize = exports.parseReferenceStorage = exports.convertClasses2Storages = exports.addStorageValues = exports.StorageType = void 0;
 const umlClass_1 = require("./umlClass");
 const associations_1 = require("./associations");
 const slotValues_1 = require("./slotValues");
+const utils_1 = require("ethers/lib/utils");
+const ethers_1 = require("ethers");
 var StorageType;
 (function (StorageType) {
-    StorageType[StorageType["Contract"] = 0] = "Contract";
-    StorageType[StorageType["Struct"] = 1] = "Struct";
+    StorageType["Contract"] = "Contract";
+    StorageType["Struct"] = "Struct";
+    StorageType["Array"] = "Array";
 })(StorageType = exports.StorageType || (exports.StorageType = {}));
 let storageId = 1;
 let variableId = 1;
@@ -18,10 +21,11 @@ let variableId = 1;
  * @param storage is mutated with the storage values
  */
 const addStorageValues = async (url, contractAddress, storage, blockTag) => {
-    const slots = storage.variables.map((s) => s.fromSlot);
+    const valueVariables = storage.variables.filter((s) => !s.noValue);
+    const slots = valueVariables.map((s) => s.fromSlot);
     const values = await (0, slotValues_1.getStorageValues)(url, contractAddress, slots, blockTag);
-    storage.variables.forEach((storage, i) => {
-        storage.values = [values[i]];
+    valueVariables.forEach((storage, i) => {
+        storage.value = values[i];
     });
 };
 exports.addStorageValues = addStorageValues;
@@ -62,8 +66,9 @@ const parseVariables = (umlClass, umlClasses, variables, storages, inheritedCont
     // Recursively parse each new inherited contract
     newInheritedContracts.forEach((parent) => {
         const parentClass = (0, associations_1.findAssociatedClass)(parent, umlClass, umlClasses);
-        if (!parentClass)
-            throw Error(`Failed to find parent contract ${parent.targetUmlClassName} of ${umlClass.absolutePath}`);
+        if (!parentClass) {
+            throw Error(`Failed to find inherited contract "${parent.targetUmlClassName}" of "${umlClass.absolutePath}"`);
+        }
         // recursively parse inherited contract
         parseVariables(parentClass, umlClasses, variables, storages, inheritedContracts);
     });
@@ -72,10 +77,11 @@ const parseVariables = (umlClass, umlClasses, variables, storages, inheritedCont
         // Ignore any attributes that are constants or immutable
         if (attribute.compiled)
             return;
-        const byteSize = (0, exports.calcStorageByteSize)(attribute, umlClass, umlClasses);
-        // find any dependent structs
-        const linkedStruct = (0, exports.parseStructStorage)(attribute, umlClasses, storages);
-        const structStorageId = linkedStruct?.id;
+        const { size: byteSize, dynamic } = (0, exports.calcStorageByteSize)(attribute, umlClass, umlClasses);
+        const noValue = attribute.attributeType === umlClass_1.AttributeType.Mapping ||
+            (attribute.attributeType === umlClass_1.AttributeType.Array && !dynamic);
+        // find any dependent storage locations
+        const referenceStorage = (0, exports.parseReferenceStorage)(attribute, umlClass, umlClasses, storages);
         // Get the toSlot of the last storage item
         let lastToSlot = 0;
         let nextOffset = 0;
@@ -84,45 +90,120 @@ const parseVariables = (umlClass, umlClasses, variables, storages, inheritedCont
             lastToSlot = lastStorage.toSlot;
             nextOffset = lastStorage.byteOffset + lastStorage.byteSize;
         }
+        let newVariable;
         if (nextOffset + byteSize > 32) {
             const nextFromSlot = variables.length > 0 ? lastToSlot + 1 : 0;
-            variables.push({
+            newVariable = {
                 id: variableId++,
                 fromSlot: nextFromSlot,
                 toSlot: nextFromSlot + Math.floor((byteSize - 1) / 32),
                 byteSize,
                 byteOffset: 0,
                 type: attribute.type,
+                dynamic,
+                noValue,
                 variable: attribute.name,
                 contractName: umlClass.name,
-                structStorageId,
-                values: [],
-            });
+                referenceStorageId: referenceStorage?.id,
+            };
         }
         else {
-            variables.push({
+            newVariable = {
                 id: variableId++,
                 fromSlot: lastToSlot,
                 toSlot: lastToSlot,
                 byteSize,
                 byteOffset: nextOffset,
                 type: attribute.type,
+                dynamic,
+                noValue,
                 variable: attribute.name,
                 contractName: umlClass.name,
-                structStorageId,
-                values: [],
-            });
+                referenceStorageId: referenceStorage?.id,
+            };
+        }
+        if (referenceStorage) {
+            if (!newVariable.dynamic) {
+                (0, exports.offsetStorageSlots)(referenceStorage, newVariable.fromSlot, storages);
+            }
+            else if (attribute.attributeType === umlClass_1.AttributeType.Array) {
+                referenceStorage.slotKey = (0, exports.calcSlotKey)(newVariable);
+            }
         }
+        variables.push(newVariable);
     });
     return variables;
 };
-const parseStructStorage = (attribute, otherClasses, storages) => {
-    if (attribute.attributeType === umlClass_1.AttributeType.UserDefined) {
-        // Have we already created the storage?
-        const existingStorage = storages.find((dep) => dep.name === attribute.type);
-        if (existingStorage) {
-            return existingStorage;
+const parseReferenceStorage = (attribute, umlClass, otherClasses, storages) => {
+    if (attribute.attributeType === umlClass_1.AttributeType.Array) {
+        // storage is dynamic if the attribute type ends in []
+        const result = attribute.type.match(/\[(\w*)]$/);
+        const dynamic = result[1] === '';
+        const arrayLength = !dynamic
+            ? (0, exports.findDimensionLength)(umlClass, result[1])
+            : undefined;
+        // get the type of the array items. eg
+        // address[][4][2] will have base type address[][4]
+        const baseType = attribute.type.substring(0, attribute.type.lastIndexOf('['));
+        let baseAttributeType;
+        if ((0, exports.isElementary)(baseType)) {
+            baseAttributeType = umlClass_1.AttributeType.Elementary;
         }
+        else if (baseType[baseType.length - 1] === ']') {
+            baseAttributeType = umlClass_1.AttributeType.Array;
+        }
+        else {
+            baseAttributeType = umlClass_1.AttributeType.UserDefined;
+        }
+        const baseAttribute = {
+            visibility: attribute.visibility,
+            name: baseType,
+            type: baseType,
+            attributeType: baseAttributeType,
+        };
+        const { size: arrayItemSize } = (0, exports.calcStorageByteSize)(baseAttribute, umlClass, otherClasses);
+        const firstVariable = {
+            id: variableId++,
+            fromSlot: 0,
+            toSlot: Math.floor((arrayItemSize - 1) / 32),
+            byteSize: arrayItemSize,
+            byteOffset: 0,
+            type: baseType,
+            dynamic,
+            noValue: false,
+        };
+        const variables = [firstVariable];
+        if (arrayLength > 1) {
+            for (let i = 1; i < arrayLength; i++) {
+                variables.push({
+                    id: variableId++,
+                    fromSlot: Math.floor((i * arrayItemSize) / 32),
+                    toSlot: Math.floor(((i + 1) * arrayItemSize - 1) / 32),
+                    byteSize: arrayItemSize,
+                    byteOffset: (i * arrayItemSize) % 32,
+                    type: baseType,
+                    dynamic,
+                    noValue: false,
+                });
+            }
+        }
+        // recursively add storage
+        if (baseAttributeType !== umlClass_1.AttributeType.Elementary) {
+            const referenceStorage = (0, exports.parseReferenceStorage)(baseAttribute, umlClass, otherClasses, storages);
+            firstVariable.referenceStorageId = referenceStorage?.id;
+        }
+        const newStorage = {
+            id: storageId++,
+            name: `${attribute.type}: ${attribute.name}`,
+            type: StorageType.Array,
+            arrayDynamic: dynamic,
+            arrayLength,
+            variables,
+        };
+        storages.push(newStorage);
+        return newStorage;
+    }
+    if (attribute.attributeType === umlClass_1.AttributeType.UserDefined) {
         // Is the user defined type linked to another Contract, Struct or Enum?
         const dependentClass = otherClasses.find(({ name }) => {
             return (name === attribute.type || name === attribute.type.split('.')[1]);
@@ -143,20 +224,13 @@ const parseStructStorage = (attribute, otherClasses, storages) => {
         }
         return undefined;
     }
-    if (attribute.attributeType === umlClass_1.AttributeType.Mapping ||
-        attribute.attributeType === umlClass_1.AttributeType.Array) {
-        // get the UserDefined type from the mapping or array
+    if (attribute.attributeType === umlClass_1.AttributeType.Mapping) {
+        // get the UserDefined type from the mapping
         // note the mapping could be an array of Structs
         // Could also be a mapping of a mapping
-        const result = attribute.attributeType === umlClass_1.AttributeType.Mapping
-            ? attribute.type.match(/=\\>((?!mapping)\w*)[\\[]/)
-            : attribute.type.match(/(\w+)\[/);
+        const result = attribute.type.match(/=\\>((?!mapping)\w*)[\\[]/);
+        // If mapping of user defined type
         if (result !== null && result[1] && !(0, exports.isElementary)(result[1])) {
-            // Have we already created the storage?
-            const existingStorage = storages.find(({ name }) => name === result[1] || name === result[1].split('.')[1]);
-            if (existingStorage) {
-                return existingStorage;
-            }
             // Find UserDefined type
             const typeClass = otherClasses.find(({ name }) => name === result[1] || name === result[1].split('.')[1]);
             if (!typeClass) {
@@ -178,66 +252,78 @@ const parseStructStorage = (attribute, otherClasses, storages) => {
     }
     return undefined;
 };
-exports.parseStructStorage = parseStructStorage;
+exports.parseReferenceStorage = parseReferenceStorage;
 // Calculates the storage size of an attribute in bytes
 const calcStorageByteSize = (attribute, umlClass, otherClasses) => {
     if (attribute.attributeType === umlClass_1.AttributeType.Mapping ||
         attribute.attributeType === umlClass_1.AttributeType.Function) {
-        return 32;
+        return { size: 32, dynamic: true };
     }
     if (attribute.attributeType === umlClass_1.AttributeType.Array) {
-        // All array dimensions must be fixed. eg [2][3][8].
-        const result = attribute.type.match(/(\w+)(\[([\w][\w]*)\])+$/);
-        // The above will not match any dynamic array dimensions, eg [],
-        // as there needs to be one or more [0-9]+ in the square brackets
-        if (result === null) {
-            // Any dynamic array dimension means the whole array is dynamic
-            // so only takes 32 bytes (1 slot)
-            return 32;
+        // Fixed sized arrays are read from right to left until there is a dynamic dimension
+        // eg address[][3][2] is a fixed size array that uses 6 slots.
+        // while address [2][] is a dynamic sized array.
+        const arrayDimensions = attribute.type.match(/\[\w*]/g);
+        // Remove first [ and last ] from each arrayDimensions
+        const dimensionsStr = arrayDimensions.map((a) => a.slice(1, -1));
+        // fixed-sized arrays are read from right to left so reverse the dimensions
+        const dimensionsStrReversed = dimensionsStr.reverse();
+        // read fixed-size dimensions until we get a dynamic array with no dimension
+        let dimension = dimensionsStrReversed.shift();
+        const fixedDimensions = [];
+        while (dimension && dimension !== '') {
+            const dimensionNum = (0, exports.findDimensionLength)(umlClass, dimension);
+            fixedDimensions.push(dimensionNum);
+            // read the next dimension for the next loop
+            dimension = dimensionsStrReversed.shift();
+        }
+        // If the first dimension is dynamic, ie []
+        if (fixedDimensions.length === 0) {
+            // dynamic arrays start at the keccak256 of the slot number
+            // the array length is stored in the 32 byte slot
+            return { size: 32, dynamic: true };
         }
-        // All array dimensions are fixes so we now need to multiply all the dimensions
-        // to get a total number of array elements
-        const arrayDimensions = attribute.type.match(/\[\w+/g);
-        const dimensionsStr = arrayDimensions.map((d) => d.slice(1));
-        const dimensions = dimensionsStr.map((dimension) => {
-            const dimensionNum = parseInt(dimension);
-            if (!isNaN(dimensionNum))
-                return dimensionNum;
-            // Try and size array dimension from declared constants
-            const constant = umlClass.constants.find((constant) => constant.name === dimension);
-            if (constant) {
-                return constant.value;
-            }
-            throw Error(`Could not size fixed sized array with dimension "${dimension}"`);
-        });
         let elementSize;
+        const type = attribute.type.substring(0, attribute.type.indexOf('['));
         // If a fixed sized array
-        if ((0, exports.isElementary)(result[1])) {
+        if ((0, exports.isElementary)(type)) {
             const elementAttribute = {
                 attributeType: umlClass_1.AttributeType.Elementary,
-                type: result[1],
+                type,
                 name: 'element',
             };
-            elementSize = (0, exports.calcStorageByteSize)(elementAttribute, umlClass, otherClasses);
+            ({ size: elementSize } = (0, exports.calcStorageByteSize)(elementAttribute, umlClass, otherClasses));
         }
         else {
             const elementAttribute = {
                 attributeType: umlClass_1.AttributeType.UserDefined,
-                type: result[1],
+                type,
                 name: 'userDefined',
             };
-            elementSize = (0, exports.calcStorageByteSize)(elementAttribute, umlClass, otherClasses);
+            ({ size: elementSize } = (0, exports.calcStorageByteSize)(elementAttribute, umlClass, otherClasses));
         }
         // Anything over 16 bytes, like an address, will take a whole 32 byte slot
         if (elementSize > 16 && elementSize < 32) {
             elementSize = 32;
         }
-        const firstDimensionBytes = elementSize * dimensions[0];
-        const firstDimensionSlotBytes = Math.ceil(firstDimensionBytes / 32) * 32;
-        const remainingElements = dimensions
-            .slice(1)
+        // If multi dimension, then the first element is 32 bytes
+        if (fixedDimensions.length < arrayDimensions.length) {
+            const totalDimensions = fixedDimensions.reduce((total, dimension) => total * dimension, 1);
+            return {
+                size: 32 * totalDimensions,
+                dynamic: false,
+            };
+        }
+        const lastItem = fixedDimensions.length - 1;
+        const lastDimensionBytes = elementSize * fixedDimensions[lastItem];
+        const lastDimensionSlotBytes = Math.ceil(lastDimensionBytes / 32) * 32;
+        const remainingDimensions = fixedDimensions
+            .slice(0, lastItem)
             .reduce((total, dimension) => total * dimension, 1);
-        return firstDimensionSlotBytes * remainingElements;
+        return {
+            size: lastDimensionSlotBytes * remainingDimensions,
+            dynamic: false,
+        };
     }
     // If a Struct or Enum
     if (attribute.attributeType === umlClass_1.AttributeType.UserDefined) {
@@ -250,12 +336,12 @@ const calcStorageByteSize = (attribute, umlClass, otherClasses) => {
         }
         switch (attributeClass.stereotype) {
             case umlClass_1.ClassStereotype.Enum:
-                return 1;
+                return { size: 1, dynamic: false };
             case umlClass_1.ClassStereotype.Contract:
             case umlClass_1.ClassStereotype.Abstract:
             case umlClass_1.ClassStereotype.Interface:
             case umlClass_1.ClassStereotype.Library:
-                return 20;
+                return { size: 20, dynamic: false };
             case umlClass_1.ClassStereotype.Struct:
                 let structByteSize = 0;
                 attributeClass.attributes.forEach((structAttribute) => {
@@ -280,7 +366,7 @@ const calcStorageByteSize = (attribute, umlClass, otherClasses) => {
                             structByteSize = Math.ceil(structByteSize / 32) * 32;
                         }
                     }
-                    const attributeSize = (0, exports.calcStorageByteSize)(structAttribute, umlClass, otherClasses);
+                    const { size: attributeSize } = (0, exports.calcStorageByteSize)(structAttribute, umlClass, otherClasses);
                     // check if attribute will fit into the remaining slot
                     const endCurrentSlot = Math.ceil(structByteSize / 32) * 32;
                     const spaceLeftInSlot = endCurrentSlot - structByteSize;
@@ -292,24 +378,27 @@ const calcStorageByteSize = (attribute, umlClass, otherClasses) => {
                     }
                 });
                 // structs take whole 32 byte slots so round up to the nearest 32 sized slots
-                return Math.ceil(structByteSize / 32) * 32;
+                return {
+                    size: Math.ceil(structByteSize / 32) * 32,
+                    dynamic: false,
+                };
             default:
-                return 32;
+                return { size: 32, dynamic: false };
         }
     }
     if (attribute.attributeType === umlClass_1.AttributeType.Elementary) {
         switch (attribute.type) {
             case 'bool':
-                return 1;
+                return { size: 1, dynamic: false };
             case 'address':
-                return 20;
+                return { size: 20, dynamic: false };
             case 'string':
             case 'bytes':
             case 'uint':
             case 'int':
             case 'ufixed':
             case 'fixed':
-                return 32;
+                return { size: 32, dynamic: false };
             default:
                 const result = attribute.type.match(/[u]*(int|fixed|bytes)([0-9]+)/);
                 if (result === null || !result[2]) {
@@ -317,12 +406,12 @@ const calcStorageByteSize = (attribute, umlClass, otherClasses) => {
                 }
                 // If bytes
                 if (result[1] === 'bytes') {
-                    return parseInt(result[2]);
+                    return { size: parseInt(result[2]), dynamic: false };
                 }
                 // TODO need to handle fixed types when they are supported
                 // If an int
                 const bitSize = parseInt(result[2]);
-                return bitSize / 8;
+                return { size: bitSize / 8, dynamic: false };
         }
     }
     throw new Error(`Failed to calc bytes size of attribute with name "${attribute.name}" and type ${attribute.type}`);
@@ -345,4 +434,44 @@ const isElementary = (type) => {
     }
 };
 exports.isElementary = isElementary;
+const calcSlotKey = (variable) => {
+    if (variable.dynamic) {
+        return (0, utils_1.keccak256)((0, utils_1.toUtf8Bytes)(ethers_1.BigNumber.from(variable.fromSlot).toHexString()));
+    }
+    return ethers_1.BigNumber.from(variable.fromSlot).toHexString();
+};
+exports.calcSlotKey = calcSlotKey;
+// recursively offset the slots numbers of a storage item
+const offsetStorageSlots = (storage, slots, storages) => {
+    storage.variables.forEach((variable) => {
+        variable.fromSlot += slots;
+        variable.toSlot += slots;
+        if (variable.referenceStorageId) {
+            // recursively offset the referenced storage
+            const referenceStorage = storages.find((s) => s.id === variable.referenceStorageId);
+            if (!referenceStorage.arrayDynamic) {
+                (0, exports.offsetStorageSlots)(referenceStorage, slots, storages);
+            }
+            else {
+                referenceStorage.slotKey = (0, exports.calcSlotKey)(variable);
+            }
+        }
+    });
+};
+exports.offsetStorageSlots = offsetStorageSlots;
+const findDimensionLength = (umlClass, dimension) => {
+    const dimensionNum = parseInt(dimension);
+    if (Number.isInteger(dimensionNum)) {
+        return dimensionNum;
+    }
+    else {
+        // Try and size array dimension from declared constants
+        const constant = umlClass.constants.find((constant) => constant.name === dimension);
+        if (!constant) {
+            throw Error(`Could not size fixed sized array with dimension "${dimension}"`);
+        }
+        return constant.value;
+    }
+};
+exports.findDimensionLength = findDimensionLength;
 //# sourceMappingURL=converterClasses2Storage.js.map

package/lib/converterStorage2Dot.d.ts

@@ -1,3 +1,7 @@
 import { Storage } from './converterClasses2Storage';
-export declare const convertStorages2Dot: (storages: Storage[]) => string;
-export declare function convertStorage2Dot(storage: Storage, dotString: string): string;
+export declare const convertStorages2Dot: (storages: Storage[], options: {
+    data: boolean;
+}) => string;
+export declare function convertStorage2Dot(storage: Storage, dotString: string, options: {
+    data: boolean;
+}): string;

package/lib/converterStorage2Dot.js

@@ -3,22 +3,22 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.convertStorage2Dot = exports.convertStorages2Dot = void 0;
 const converterClasses2Storage_1 = require("./converterClasses2Storage");
 const debug = require('debug')('sol2uml');
-const convertStorages2Dot = (storages) => {
+const convertStorages2Dot = (storages, options) => {
     let dotString = `
 digraph StorageDiagram {
 rankdir=LR
 color=black
 arrowhead=open
-node [shape=record, style=filled, fillcolor=gray95]`;
+node [shape=record, style=filled, fillcolor=gray95 fontname="Courier New"]`;
     // process contract and the struct storages
     storages.forEach((storage) => {
-        dotString = convertStorage2Dot(storage, dotString);
+        dotString = convertStorage2Dot(storage, dotString, options);
     });
     // link contract and structs to structs
     storages.forEach((slot) => {
         slot.variables.forEach((storage) => {
-            if (storage.structStorageId) {
-                dotString += `\n ${slot.id}:${storage.id} -> ${storage.structStorageId}`;
+            if (storage.referenceStorageId) {
+                dotString += `\n ${slot.id}:${storage.id} -> ${storage.referenceStorageId}`;
             }
         });
     });
@@ -28,10 +28,10 @@ node [shape=record, style=filled, fillcolor=gray95]`;
     return dotString;
 };
 exports.convertStorages2Dot = convertStorages2Dot;
-function convertStorage2Dot(storage, dotString) {
-    const steorotype = storage.type === converterClasses2Storage_1.StorageType.Struct ? 'Struct' : 'Contract';
+function convertStorage2Dot(storage, dotString, options) {
     // write storage header with name and optional address
-    dotString += `\n${storage.id} [label="${storage.name} \\<\\<${steorotype}\\>\\>\\n${storage.address || ''} | {`;
+    dotString += `\n${storage.id} [label="${storage.name} \\<\\<${storage.type}\\>\\>\\n${storage.address || storage.slotKey || ''}`;
+    dotString += ' | {';
     const startingVariables = storage.variables.filter((s) => s.byteOffset === 0);
     // write slot numbers
     dotString += '{ slot';
@@ -44,10 +44,10 @@ function convertStorage2Dot(storage, dotString) {
         }
     });
     // write slot values if available
-    if (startingVariables[0]?.values[0]) {
+    if (options.data) {
         dotString += '} | {value';
         startingVariables.forEach((variable, i) => {
-            dotString += ` | ${variable.values[0]}`;
+            dotString += ` | ${variable.value || ''}`;
         });
     }
     const contractVariablePrefix = storage.type === converterClasses2Storage_1.StorageType.Contract ? '\\<inherited contract\\>.' : '';
@@ -58,6 +58,7 @@ function convertStorage2Dot(storage, dotString) {
         const slotVariables = storage.variables.filter((s) => s.fromSlot === variable.fromSlot);
         const usedBytes = slotVariables.reduce((acc, s) => acc + s.byteSize, 0);
         if (usedBytes < 32) {
+            // Create an unallocated variable for display purposes
             slotVariables.push({
                 id: 0,
                 fromSlot: variable.fromSlot,
@@ -65,9 +66,10 @@ function convertStorage2Dot(storage, dotString) {
                 byteSize: 32 - usedBytes,
                 byteOffset: usedBytes,
                 type: 'unallocated',
+                dynamic: false,
+                noValue: true,
                 contractName: variable.contractName,
                 variable: '',
-                values: [],
             });
         }
         const slotVariablesReversed = slotVariables.reverse();
@@ -88,7 +90,7 @@ function convertStorage2Dot(storage, dotString) {
 }
 exports.convertStorage2Dot = convertStorage2Dot;
 const dotVariable = (storage, contractName) => {
-    const port = storage.structStorageId !== undefined ? `<${storage.id}>` : '';
+    const port = storage.referenceStorageId !== undefined ? `<${storage.id}>` : '';
     const contractNamePrefix = storage.contractName !== contractName ? `${storage.contractName}.` : '';
     const variable = storage.variable
         ? `: ${contractNamePrefix}${storage.variable}`

package/lib/filterClasses.d.ts

@@ -1,6 +1,7 @@
 import { WeightedDiGraph } from 'js-graph-algorithms';
 import { UmlClass } from './umlClass';
 export declare const classesConnectedToBaseContracts: (umlClasses: UmlClass[], baseContractNames: string[], depth?: number) => UmlClass[];
-export declare const classesConnectedToBaseContract: (umlClasses: UmlClass[], baseContractName: string, graph: WeightedDiGraph, depth?: number) => {
+export declare const classesConnectedToBaseContract: (umlClasses: UmlClass[], baseContractName: string, weightedDirectedGraph: WeightedDiGraph, depth?: number) => {
     [contractName: string]: UmlClass;
 };
+export declare const topologicalSortClasses: (umlClasses: UmlClass[]) => UmlClass[];

package/lib/filterClasses.js

@@ -1,21 +1,21 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.classesConnectedToBaseContract = exports.classesConnectedToBaseContracts = void 0;
+exports.topologicalSortClasses = exports.classesConnectedToBaseContract = exports.classesConnectedToBaseContracts = void 0;
 const js_graph_algorithms_1 = require("js-graph-algorithms");
 const associations_1 = require("./associations");
 const classesConnectedToBaseContracts = (umlClasses, baseContractNames, depth) => {
     let filteredUmlClasses = {};
-    const graph = loadGraph(umlClasses);
+    const weightedDirectedGraph = loadWeightedDirectedGraph(umlClasses);
     for (const baseContractName of baseContractNames) {
         filteredUmlClasses = {
             ...filteredUmlClasses,
-            ...(0, exports.classesConnectedToBaseContract)(umlClasses, baseContractName, graph, depth),
+            ...(0, exports.classesConnectedToBaseContract)(umlClasses, baseContractName, weightedDirectedGraph, depth),
         };
     }
     return Object.values(filteredUmlClasses);
 };
 exports.classesConnectedToBaseContracts = classesConnectedToBaseContracts;
-const classesConnectedToBaseContract = (umlClasses, baseContractName, graph, depth = 1000) => {
+const classesConnectedToBaseContract = (umlClasses, baseContractName, weightedDirectedGraph, depth = 1000) => {
     // Find the base UML Class from the base contract name
     const baseUmlClass = umlClasses.find(({ name }) => {
         return name === baseContractName;
@@ -23,7 +23,7 @@ const classesConnectedToBaseContract = (umlClasses, baseContractName, graph, dep
     if (!baseUmlClass) {
         throw Error(`Failed to find base contract with name "${baseContractName}"`);
     }
-    const dfs = new js_graph_algorithms_1.Dijkstra(graph, baseUmlClass.id);
+    const dfs = new js_graph_algorithms_1.Dijkstra(weightedDirectedGraph, baseUmlClass.id);
     // Get all the UML Classes that are connected to the base contract
     const filteredUmlClasses = {};
     for (const umlClass of umlClasses) {
@@ -34,8 +34,8 @@ const classesConnectedToBaseContract = (umlClasses, baseContractName, graph, dep
     return filteredUmlClasses;
 };
 exports.classesConnectedToBaseContract = classesConnectedToBaseContract;
-function loadGraph(umlClasses) {
-    const graph = new js_graph_algorithms_1.WeightedDiGraph(umlClasses.length); // 6 is the number vertices in the graph
+function loadWeightedDirectedGraph(umlClasses) {
+    const weightedDirectedGraph = new js_graph_algorithms_1.WeightedDiGraph(umlClasses.length); // the number vertices in the graph
     for (const sourceUmlClass of umlClasses) {
         for (const association of Object.values(sourceUmlClass.associations)) {
             // Find the first UML Class that matches the target class name
@@ -43,9 +43,34 @@ function loadGraph(umlClasses) {
             if (!targetUmlClass) {
                 continue;
             }
-            graph.addEdge(new js_graph_algorithms_1.Edge(sourceUmlClass.id, targetUmlClass.id, 1));
+            weightedDirectedGraph.addEdge(new js_graph_algorithms_1.Edge(sourceUmlClass.id, targetUmlClass.id, 1));
         }
     }
-    return graph;
+    return weightedDirectedGraph;
 }
+const topologicalSortClasses = (umlClasses) => {
+    const directedAcyclicGraph = loadDirectedAcyclicGraph(umlClasses);
+    const topologicalSort = new js_graph_algorithms_1.TopologicalSort(directedAcyclicGraph);
+    // Topological sort the class ids
+    const sortedUmlClassIds = topologicalSort.order().reverse();
+    const sortedUmlClasses = sortedUmlClassIds.map((umlClassId) => 
+    // Lookup the UmlClass for each class id
+    umlClasses.find((umlClass) => umlClass.id === umlClassId));
+    return sortedUmlClasses;
+};
+exports.topologicalSortClasses = topologicalSortClasses;
+const loadDirectedAcyclicGraph = (umlClasses) => {
+    const directedAcyclicGraph = new js_graph_algorithms_1.DiGraph(umlClasses.length); // the number vertices in the graph
+    for (const sourceUmlClass of umlClasses) {
+        for (const association of Object.values(sourceUmlClass.associations)) {
+            // Find the first UML Class that matches the target class name
+            const targetUmlClass = (0, associations_1.findAssociatedClass)(association, sourceUmlClass, umlClasses);
+            if (!targetUmlClass) {
+                continue;
+            }
+            directedAcyclicGraph.addEdge(sourceUmlClass.id, targetUmlClass.id);
+        }
+    }
+    return directedAcyclicGraph;
+};
 //# sourceMappingURL=filterClasses.js.map

package/lib/parserEtherscan.js

@@ -7,6 +7,8 @@ exports.EtherscanParser = void 0;
 const axios_1 = __importDefault(require("axios"));
 const parser_1 = require("@solidity-parser/parser");
 const converterAST2Classes_1 = require("./converterAST2Classes");
+const filterClasses_1 = require("./filterClasses");
+const debug = require('debug')('sol2uml');
 const networks = [
     'mainnet',
     'ropsten',
@@ -69,9 +71,39 @@ class EtherscanParser {
      */
     async getSolidityCode(contractAddress) {
         const { files, contractName } = await this.getSourceCode(contractAddress);
+        // Parse the UmlClasses from the Solidity code in each file
+        let umlClasses = [];
+        for (const file of files) {
+            const node = await this.parseSourceCode(file.code);
+            const umlClass = (0, converterAST2Classes_1.convertAST2UmlClasses)(node, file.filename);
+            umlClasses = umlClasses.concat(umlClass);
+        }
+        // Sort the classes so dependent code is first
+        const topologicalSortedClasses = (0, filterClasses_1.topologicalSortClasses)(umlClasses);
+        // Get a list of filenames the classes are in
+        const sortedFilenames = topologicalSortedClasses.map((umlClass) => umlClass.relativePath);
+        // Remove duplicate filenames from the list
+        const dependentFilenames = [...new Set(sortedFilenames)];
+        // find any files that didn't have dependencies found
+        const nonDependentFiles = files.filter((f) => !dependentFilenames.includes(f.filename));
+        const nonDependentFilenames = nonDependentFiles.map((f) => f.filename);
         let solidityCode = '';
-        files.forEach((file) => {
-            solidityCode += file.code;
+        // output non dependent code before the dependent files just in case sol2uml missed some dependencies
+        const filenames = [...nonDependentFilenames, ...dependentFilenames];
+        // For each filename
+        filenames.forEach((filename) => {
+            // Lookup the file that contains the Solidity code
+            const file = files.find((f) => f.filename === filename);
+            if (!file)
+                throw Error(`Failed to find file with filename "${filename}"`);
+            // comment out any import statements
+            // match whitespace before import
+            // and characters after import up to ;
+            // replace all in file and match across multiple lines
+            const removedImports = file.code.replace(/(\s)(import.*;)/gm, '$1/* $2 */');
+            // Rename SPDX-License-Identifier to SPDX--License-Identifier so the merged file will compile
+            const removedSPDX = removedImports.replace(/SPDX-/, 'SPDX--');
+            solidityCode += removedSPDX;
         });
         return {
             solidityCode,
@@ -99,6 +131,7 @@ class EtherscanParser {
     async getSourceCode(contractAddress) {
         const description = `get verified source code for address ${contractAddress} from Etherscan API.`;
         try {
+            debug(`About to get Solidity source code for ${contractAddress} from ${this.url}`);
             const response = await axios_1.default.get(this.url, {
                 params: {
                     module: 'contract',

package/lib/slotValues.js

@@ -41,7 +41,7 @@ const getStorageValues = async (url, contractAddress, slots, blockTag = 'latest'
         }
         const responseData = response.data;
         const sortedResponses = responseData.sort((a, b) => bignumber_1.BigNumber.from(a.id).gt(b.id) ? 1 : -1);
-        return sortedResponses.map((data) => data.result);
+        return sortedResponses.map((data) => '0x' + data.result.toUpperCase().slice(2));
     }
     catch (err) {
         throw new Error(`Failed to get ${slots.length} storage values for ${contractAddress} from ${url}`, { cause: err });

package/lib/sol2uml.js

@@ -11,14 +11,16 @@ const converterStorage2Dot_1 = require("./converterStorage2Dot");
 const regEx_1 = require("./utils/regEx");
 const writerFiles_1 = require("./writerFiles");
 const program = new commander_1.Command();
+program.version(require('../package.json').version);
 const debugControl = require('debug');
 const debug = require('debug')('sol2uml');
 program
     .usage(`[subcommand] <options>
-The three subcommands:
-* class:    Generates a UML class diagram from Solidity source code. default
+
+sol2uml comes with three subcommands:
+* class:    Generates a UML class diagram from Solidity source code. (default)
 * storage:  Generates a diagram of a contract's storage slots.
-* flatten:  Pulls verified source files from a Blockchain explorer into one, flat, local Solidity file.
+* flatten:  Merges verified source files from a Blockchain explorer into one local file.
 
 The Solidity code can be pulled from verified source code on Blockchain explorers like Etherscan or from local Solidity files.`)
     .addOption(new commander_1.Option('-sf, --subfolders <value>', 'number of subfolders that will be recursively searched for Solidity files.').default('-1', 'all'))
@@ -39,8 +41,9 @@ The Solidity code can be pulled from verified source code on Blockchain explorer
     'goerli',
     'sepolia',
 ])
-    .default('mainnet'))
-    .option('-k, --apiKey <key>', 'Etherscan, Polygonscan, BscScan or Arbiscan API key')
+    .default('mainnet')
+    .env('ETH_NETWORK'))
+    .addOption(new commander_1.Option('-k, --apiKey <key>', 'Etherscan, Polygonscan, BscScan or Arbiscan API key').env('SCAN_API_KEY'))
     .option('-v, --verbose', 'run with debugging statements', false);
 program
     .command('class', { isDefault: true })
@@ -75,23 +78,27 @@ If an Ethereum address with a 0x prefix is passed, the verified source code from
             ...command.parent._optionValues,
             ...options,
         };
-        const { umlClasses, contractName } = await (0, parserGeneral_1.parserUmlClasses)(fileFolderAddress, combinedOptions);
+        let { umlClasses, contractName } = await (0, parserGeneral_1.parserUmlClasses)(fileFolderAddress, combinedOptions);
         let filteredUmlClasses = umlClasses;
         if (options.baseContractNames) {
             const baseContractNames = options.baseContractNames.split(',');
             filteredUmlClasses = (0, filterClasses_1.classesConnectedToBaseContracts)(umlClasses, baseContractNames, options.depth);
+            contractName = baseContractNames[0];
         }
         const dotString = (0, converterClasses2Dot_1.convertUmlClasses2Dot)(filteredUmlClasses, combinedOptions.clusterFolders, combinedOptions);
-        await (0, writerFiles_1.writeOutputFiles)(dotString, fileFolderAddress, contractName, combinedOptions.outputFormat, combinedOptions.outputFileName);
+        await (0, writerFiles_1.writeOutputFiles)(dotString, fileFolderAddress, contractName || 'classDiagram', combinedOptions.outputFormat, combinedOptions.outputFileName);
         debug(`Finished generating UML`);
     }
     catch (err) {
-        console.error(`Failed to generate UML diagram ${err}`);
+        console.error(`Failed to generate UML diagram\n${err.stack}`);
     }
 });
 program
     .command('storage')
-    .description("Visually display a contract's storage slots.\n\nWARNING: sol2uml does not use the Solidity compiler so may differ with solc. A known example is fixed-sized arrays declared with an expression will fail to be sized.")
+    .description("Visually display a contract's storage slots.")
+    .usage(`<fileFolderAddress> [options]
+
+WARNING: sol2uml does not use the Solidity compiler so may differ with solc. A known example is fixed-sized arrays declared with an expression will fail to be sized.`)
     .argument('<fileFolderAddress>', 'file name, base folder or contract address')
     .option('-c, --contract <name>', 'Contract name in local Solidity files. Not needed when using an address as the first argument as the contract name can be derived from Etherscan.')
     .option('-d, --data', 'Gets the values in the storage slots from an Ethereum node.', false)
@@ -132,28 +139,39 @@ program
                 throw Error(`Could not find the "${contractName}" contract in list of parsed storages`);
             await (0, converterClasses2Storage_1.addStorageValues)(combinedOptions.url, storageAddress, storage, combinedOptions.blockNumber);
         }
-        const dotString = (0, converterStorage2Dot_1.convertStorages2Dot)(storages);
-        await (0, writerFiles_1.writeOutputFiles)(dotString, fileFolderAddress, contractName, combinedOptions.outputFormat, combinedOptions.outputFileName);
+        const dotString = (0, converterStorage2Dot_1.convertStorages2Dot)(storages, combinedOptions);
+        await (0, writerFiles_1.writeOutputFiles)(dotString, fileFolderAddress, contractName || 'storageDiagram', combinedOptions.outputFormat, combinedOptions.outputFileName);
     }
     catch (err) {
-        console.error(`Failed to generate storage diagram ${err}`);
+        console.error(`Failed to generate storage diagram.\n${err.stack}`);
     }
 });
 program
     .command('flatten')
-    .description('get all verified source code for a contract from the Blockchain explorer into one local file')
-    .argument('<contractAddress>', 'Contract address')
+    .description('Merges verified source files for a contract from a Blockchain explorer into one local file.')
+    .usage(`<contractAddress> [options]
+
+In order for the merged code to compile, the following is done:
+1. File imports are commented out.
+2. "SPDX-License-Identifier" is renamed to "SPDX--License-Identifier".
+3. Contract dependencies are analysed so the files are merged in an order that will compile.`)
+    .argument('<contractAddress>', 'Contract address in hexadecimal format with a 0x prefix.')
     .action(async (contractAddress, options, command) => {
-    debug(`About to flatten ${contractAddress}`);
-    const combinedOptions = {
-        ...command.parent._optionValues,
-        ...options,
-    };
-    const etherscanParser = new parserEtherscan_1.EtherscanParser(combinedOptions.apiKey, combinedOptions.network);
-    const { solidityCode, contractName } = await etherscanParser.getSolidityCode(contractAddress);
-    // Write Solidity to the contract address
-    const outputFilename = combinedOptions.outputFileName || contractName;
-    await (0, writerFiles_1.writeSolidity)(solidityCode, outputFilename);
+    try {
+        debug(`About to flatten ${contractAddress}`);
+        const combinedOptions = {
+            ...command.parent._optionValues,
+            ...options,
+        };
+        const etherscanParser = new parserEtherscan_1.EtherscanParser(combinedOptions.apiKey, combinedOptions.network);
+        const { solidityCode, contractName } = await etherscanParser.getSolidityCode(contractAddress);
+        // Write Solidity to the contract address
+        const outputFilename = combinedOptions.outputFileName || contractName;
+        await (0, writerFiles_1.writeSolidity)(solidityCode, outputFilename);
+    }
+    catch (err) {
+        console.error(`Failed to flatten files.\n${err.stack}`);
+    }
 });
 program.on('option:verbose', () => {
     debugControl.enable('sol2uml');

package/lib/umlClass.d.ts

@@ -88,7 +88,7 @@ export declare class UmlClass implements ClassProperties {
     name: string;
     absolutePath: string;
     relativePath: string;
-    imports?: Import[];
+    imports: Import[];
     stereotype?: ClassStereotype;
     constants: Constants[];
     attributes: Attribute[];

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "sol2uml",
-  "version": "2.1.0",
+  "version": "2.1.3",
   "description": "Unified Modeling Language (UML) class diagram generator for Solidity contracts",
   "main": "./lib/index.js",
   "types": "./lib/index.d.ts",