From e5ad9f2c923d6e582b678f9853bb70e9e2c7395c Mon Sep 17 00:00:00 2001
From: Harlan Haskins <harlan@harlanhaskins.com>
Date: Sat, 14 Jan 2017 14:44:05 -0500
Subject: [PATCH] Added iteration attributes to Module, Global, and Function

---
 Sources/LLVM/Function.swift | 12 ++++++++
 Sources/LLVM/Global.swift   | 12 ++++++++
 Sources/LLVM/Module.swift   | 55 ++++++++++++++++++++++++++++++++++++-
 3 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/Sources/LLVM/Function.swift b/Sources/LLVM/Function.swift
index 1db09ff1..2c3b71c3 100644
--- a/Sources/LLVM/Function.swift
+++ b/Sources/LLVM/Function.swift
@@ -50,6 +50,18 @@ public class Function: IRValue {
     }
   }
 
+  /// Retrieves the previous function in the module, if there is one.
+  public func previous() -> Function? {
+    guard let previous = LLVMGetPreviousFunction(llvm) else { return nil }
+    return Function(llvm: previous)
+  }
+
+  /// Retrieves the next function in the module, if there is one.
+  public func next() -> Function? {
+    guard let next = LLVMGetNextFunction(llvm) else { return nil }
+    return Function(llvm: next)
+  }
+
   /// Retrieves a parameter at the given index, if it exists.
   ///
   /// - parameter index: The index of the parameter to retrieve.
diff --git a/Sources/LLVM/Global.swift b/Sources/LLVM/Global.swift
index 97ef4e03..2e7224d4 100644
--- a/Sources/LLVM/Global.swift
+++ b/Sources/LLVM/Global.swift
@@ -33,6 +33,18 @@ public struct Global: IRValue {
     set { LLVMSetThreadLocal(asLLVM(), newValue.llvm) }
   }
 
+  /// Retrieves the previous global in the module, if there is one.
+  public func previous() -> Global? {
+    guard let previous = LLVMGetPreviousGlobal(llvm) else { return nil }
+    return Global(llvm: previous)
+  }
+
+  /// Retrieves the next global in the module, if there is one.
+  public func next() -> Global? {
+    guard let next = LLVMGetNextGlobal(llvm) else { return nil }
+    return Global(llvm: next)
+  }
+
   /// Deletes the global variable from its containing module.
   /// - note: This does not remove references to this global from the
   ///         module. Ensure you have removed all insructions that reference
diff --git a/Sources/LLVM/Module.swift b/Sources/LLVM/Module.swift
index ecc17303..9df424c9 100644
--- a/Sources/LLVM/Module.swift
+++ b/Sources/LLVM/Module.swift
@@ -41,7 +41,7 @@ public enum ModuleError: Error, CustomStringConvertible {
 /// A `Module` represents the top-level structure of an LLVM program. An LLVM
 /// module is effectively a translation unit or a collection of translation
 /// units merged together.
-public final class Module {
+public final class Module: CustomStringConvertible {
   internal let llvm: LLVMModuleRef
 
   /// Creates a `Module` with the given name.
@@ -123,11 +123,64 @@ public final class Module {
     }
   }
 
+  /// Retrieves the sequence of functions that make up this module.
+  public var functions: AnySequence<Function> {
+    var current = firstFunction
+    return AnySequence<Function> {
+      return AnyIterator<Function> {
+        defer { current = current?.next() }
+        return current
+      }
+    }
+  }
+
+  /// Retrieves the first function in this module, if there are any functions.
+  public var firstFunction: Function? {
+    guard let fn = LLVMGetFirstFunction(llvm) else { return nil }
+    return Function(llvm: fn)
+  }
+
+  /// Retrieves the last function in this module, if there are any functions.
+  public var lastFunction: Function? {
+    guard let fn = LLVMGetLastFunction(llvm) else { return nil }
+    return Function(llvm: fn)
+  }
+
+  /// Retrieves the first global in this module, if there are any globals.
+  public var firstGlobal: Global? {
+    guard let fn = LLVMGetFirstGlobal(llvm) else { return nil }
+    return Global(llvm: fn)
+  }
+
+  /// Retrieves the last global in this module, if there are any globals.
+  public var lastGlobal: Global? {
+    guard let fn = LLVMGetLastGlobal(llvm) else { return nil }
+    return Global(llvm: fn)
+  }
+
+  /// Retrieves the sequence of functions that make up this module.
+  public var globals: AnySequence<Global> {
+    var current = firstGlobal
+    return AnySequence<Global> {
+      return AnyIterator<Global> {
+        defer { current = current?.next() }
+        return current
+      }
+    }
+  }
+
   /// Dump a representation of this module to stderr.
   public func dump() {
     LLVMDumpModule(llvm)
   }
 
+  /// The full text IR of this module
+  public var description: String {
+    let cStr = LLVMPrintModuleToString(llvm)!
+    defer { LLVMDisposeMessage(cStr) }
+    return String(cString: cStr)
+  }
+
   deinit {
     LLVMDisposeModule(llvm)
   }