Linter for Calldata Block Iterator

2 files changed, 26 insertions, 13 deletions

diff --git a/packages/utils/src/abi_encoder/calldata/calldata.ts b/packages/utils/src/abi_encoder/calldata/calldata.ts
index dd9d47def..50f0f0fad 100644
--- a/packages/utils/src/abi_encoder/calldata/calldata.ts
+++ b/packages/utils/src/abi_encoder/calldata/calldata.ts
@@ -31,8 +31,7 @@ export class Calldata {
         // 1. Create a queue of subtrees by hash
         // Note that they are ordered the same as
         const iterator = new ReverseCalldataIterator(this._root);
-        let block: CalldataBlock | undefined;
-        while (block = iterator.next()) {
+        for (const block of iterator) {
             if (block instanceof CalldataBlocks.Pointer) {
                 const dependencyBlockHashBuf = block.getDependency().computeHash();
                 const dependencyBlockHash = ethUtil.bufferToHex(dependencyBlockHashBuf);
@@ -63,9 +62,8 @@ export class Calldata {
         }
 
         const iterator = new CalldataIterator(this._root);
-        let block: CalldataBlock | undefined;
         let offset = 0;
-        while (block = iterator.next()) {
+        for (const block of iterator) {
             block.setOffset(offset);
             offset += block.getSizeInBytes();
         }
@@ -102,11 +100,9 @@ export class Calldata {
         }
 
         const iterator = new CalldataIterator(this._root);
-
-        let block: CalldataBlock | undefined;
         let offset = 0;
         const functionName: string = this._root.getName();
-        while (block = iterator.next()) {
+        for (const block of iterator) {
             // Process each block 1 word at a time
             const size = block.getSizeInBytes();
             const name = block.getName();
@@ -175,8 +171,7 @@ export class Calldata {
 
         const iterator = new CalldataIterator(this._root);
         const valueBufs: Buffer[] = [selectorBuffer];
-        let block: CalldataBlock | undefined;
-        while (block = iterator.next()) {
+        for (const block of iterator) {
             valueBufs.push(block.toBuffer());
         }
 
diff --git a/packages/utils/src/abi_encoder/calldata/iterator.ts b/packages/utils/src/abi_encoder/calldata/iterator.ts
index 3e3367e10..8e2b16a5a 100644
--- a/packages/utils/src/abi_encoder/calldata/iterator.ts
+++ b/packages/utils/src/abi_encoder/calldata/iterator.ts
@@ -35,7 +35,7 @@ import { CalldataBlock } from './calldata_block';
  *            }
  * It will iterate as follows: [A, B, C, B.a, B.b, C.c]
  */
-abstract class BaseIterator {
+abstract class BaseIterator implements Iterable<CalldataBlock> {
     protected readonly _root: CalldataBlock;
     protected readonly _queue: Queue<CalldataBlock>;
 
@@ -70,7 +70,25 @@ abstract class BaseIterator {
         this._queue = BaseIterator._createQueue(root);
     }
 
-    public abstract next(): CalldataBlock | undefined;
+    public [Symbol.iterator](): { next: () => IteratorResult<CalldataBlock> } {
+        return {
+            next: () => {
+                const nextBlock = this.nextBlock();
+                if (nextBlock !== undefined) {
+                    return {
+                        value: nextBlock,
+                        done: false,
+                    };
+                }
+                return {
+                    done: true,
+                    value: new CalldataBlocks.Blob('', '', '', new Buffer('')),
+                };
+            },
+        };
+    }
+
+    public abstract nextBlock(): CalldataBlock | undefined;
 }
 
 export class CalldataIterator extends BaseIterator {
@@ -78,7 +96,7 @@ export class CalldataIterator extends BaseIterator {
         super(root);
     }
 
-    public next(): CalldataBlock | undefined {
+    public nextBlock(): CalldataBlock | undefined {
         return this._queue.popFront();
     }
 }
@@ -88,7 +106,7 @@ export class ReverseCalldataIterator extends BaseIterator {
         super(root);
     }
 
-    public next(): CalldataBlock | undefined {
+    public nextBlock(): CalldataBlock | undefined {
         return this._queue.popBack();
     }
 }