diff --git a/Sources/MockoloFramework/Models/ClassModel.swift b/Sources/MockoloFramework/Models/ClassModel.swift
index 395ebdaa..5ed68f14 100644
--- a/Sources/MockoloFramework/Models/ClassModel.swift
+++ b/Sources/MockoloFramework/Models/ClassModel.swift
@@ -46,7 +46,7 @@ final class ClassModel: Model {
         self.needInit = needInit
         self.initParams = initParams
         self.offset = offset
-        self.attribute = Set(attributes).joined(separator: " ")
+        self.attribute = Set(attributes.filter {$0.contains(String.available)}).joined(separator: " ")
         self.accessControlLevelDescription = acl.isEmpty ? "" : acl + " "
         self.typealiasWhitelist = typealiasWhitelist
     }
diff --git a/Sources/MockoloFramework/Models/MethodModel.swift b/Sources/MockoloFramework/Models/MethodModel.swift
index 2a49e52b..e91638fd 100644
--- a/Sources/MockoloFramework/Models/MethodModel.swift
+++ b/Sources/MockoloFramework/Models/MethodModel.swift
@@ -135,7 +135,7 @@ final class MethodModel: Model {
         
         let zippedParams = zip(paramDecls, comps)
         self.params = zippedParams.map { (argAst: Structure, argLabel: String) -> ParamModel in
-            ParamModel(argAst, label: argLabel, offset: argAst.offset, length: argAst.length, data: data, isInitializer: ast.isInitializer)
+            return ParamModel(argAst, label: argLabel, offset: argAst.offset, length: argAst.length, data: data, inInit: ast.isInitializer)
         }
         
         self.genericTypeParams = ast.substructures
diff --git a/Sources/MockoloFramework/Models/ParamModel.swift b/Sources/MockoloFramework/Models/ParamModel.swift
index 5b78340b..b81f1bc0 100644
--- a/Sources/MockoloFramework/Models/ParamModel.swift
+++ b/Sources/MockoloFramework/Models/ParamModel.swift
@@ -24,13 +24,13 @@ final class ParamModel: Model {
     var type: Type
     let label: String
     let isGeneric: Bool
-    let isInitializer: Bool
+    let inInit: Bool
     
     var modelType: ModelType {
         return .parameter
     }
 
-    init(label: String = "", name: String, typeName: String, isGeneric: Bool = false, isInitializer: Bool = false, offset: Int64, length: Int64) {
+    init(label: String = "", name: String, typeName: String, isGeneric: Bool = false, inInit: Bool = false, offset: Int64, length: Int64) {
         self.name = name.trimmingCharacters(in: .whitespaces)
         self.type = Type(typeName.trimmingCharacters(in: .whitespaces))
         let labelStr = label.trimmingCharacters(in: .whitespaces)
@@ -38,10 +38,10 @@ final class ParamModel: Model {
         self.offset = offset
         self.length = length
         self.isGeneric = isGeneric
-        self.isInitializer = isInitializer
+        self.inInit = inInit
     }
     
-    init(_ ast: Structure, label: String = "", offset: Int64, length: Int64, data: Data, isGeneric: Bool = false, isInitializer: Bool = false) {
+    init(_ ast: Structure, label: String = "", offset: Int64, length: Int64, data: Data, isGeneric: Bool = false, inInit: Bool = false) {
         self.name = ast.name
         self.offset = offset
         self.length = length
@@ -49,14 +49,14 @@ final class ParamModel: Model {
         let lookahead = data.toString(offset: offset + length, length: 3)
         let isVariadic = lookahead == "..."
         self.isGeneric = isGeneric
-        self.isInitializer = isInitializer
+        self.inInit = inInit
         let typeArg = isGeneric ? (ast.inheritedTypes.first ?? .unknownVal) : (isVariadic ? ast.typeName + "..." : ast.typeName)
         self.type = Type(typeArg)
         self.label = ast.name != label ? label: ""
     }
 
     var asVarDecl: String? {
-        if self.isInitializer {
+        if self.inInit {
             return applyVarTemplate(name: name, type: type)
         }
         return nil
diff --git a/Sources/MockoloFramework/Models/ParsedEntity.swift b/Sources/MockoloFramework/Models/ParsedEntity.swift
index 023dbbc0..91a236a2 100644
--- a/Sources/MockoloFramework/Models/ParsedEntity.swift
+++ b/Sources/MockoloFramework/Models/ParsedEntity.swift
@@ -29,9 +29,9 @@ struct ResolvedEntity {
     
     func model() -> Model {
         return ClassModel(identifier: key,
-                          acl: entity.acl,
+                          acl: entity.entityNode.acl,
                           attributes: attributes,
-                          offset: entity.offset,
+                          offset: entity.entityNode.offset,
                           needInit: !hasInit,
                           initParams: initVars,
                           typealiasWhitelist: typealiasWhitelist,
@@ -45,65 +45,48 @@ struct ResolvedEntityContainer {
     let imports: [(String, Data, Int64)]
 }
 
+protocol EntityNode {
+    var name: String { get }
+    var acl: String { get }
+    var attributesDescription: String { get }
+    var inheritedTypes: [String] { get }
+    var offset: Int64 { get }
+    func subContainer(overrides: [String: String]?, path: String?, data: Data?, isProcessed: Bool) -> EntityNodeSubContainer
+}
+
+final class EntityNodeSubContainer {
+    let attributes: [String]
+    let members: [Model]
+    let hasInit: Bool
+    init(attributes: [String], members: [Model], hasInit: Bool) {
+        self.attributes = attributes
+        self.members = members
+        self.hasInit = hasInit
+    }
+}
+
+
 /// Metadata for a type being mocked
 final class Entity {
     var filepath: String = ""
     var data: Data? = nil
 
-    let name: String
-    let members: [Model]
-    let offset: Int64
-    let acl: String
-    let attributes: [String]
-    let inheritedTypes: [String]
-    let hasInit: Bool
     let isAnnotated: Bool
     let overrides: [String: String]?
+    let entityNode: EntityNode
     let isProcessed: Bool
     
-    init(name: String,
+    init(entityNode: EntityNode,
          filepath: String = "",
          data: Data? = nil,
          isAnnotated: Bool,
          overrides: [String: String]?,
-         acl: String,
-         attributes: [String],
-         inheritedTypes: [String],
-         members: [Model],
-         hasInit: Bool,
-         offset: Int64,
          isProcessed: Bool) {
-        self.name = name
+        self.entityNode = entityNode
         self.filepath = filepath
         self.data = data
-        self.acl = acl
-        self.attributes = attributes
-        self.inheritedTypes = inheritedTypes
-        self.hasInit = hasInit
         self.isAnnotated = isAnnotated
         self.overrides = overrides
         self.isProcessed = isProcessed
-        self.offset = offset
-        self.members = members
-    }
-
-    
-    func subAttributes() -> [String]? {
-        if isProcessed {
-            return nil
-        }
-        return attributes.filter {$0.contains(String.available)}
-    }
-
-    static func model(for element: Structure, filepath: String, data: Data, overrides: [String: String]?, processed: Bool = false) -> Model? {
-        if element.isVariable {
-            return VariableModel(element, filepath: filepath, data: data, processed: processed)
-        } else if element.isMethod {
-            return MethodModel(element, filepath: filepath, data: data,  processed: processed)
-        } else if element.isAssociatedType {
-            return TypeAliasModel(element, filepath: filepath, data: data, overrideTypes: overrides, processed: processed)
-        }
-        
-        return nil
     }
 }
diff --git a/Sources/MockoloFramework/Operations/Generator.swift b/Sources/MockoloFramework/Operations/Generator.swift
index 8e35e6ef..a529f25f 100644
--- a/Sources/MockoloFramework/Operations/Generator.swift
+++ b/Sources/MockoloFramework/Operations/Generator.swift
@@ -68,7 +68,7 @@ public func generate(sourceDirs: [String]?,
                                  semaphore: sema,
                                  queue: mockgenQueue) { (elements, imports) in
                                     elements.forEach { element in
-                                        parentMocks[element.name] = element
+                                        parentMocks[element.entityNode.name] = element
                                     }
 
                                     for (path, modules) in imports {
@@ -94,9 +94,9 @@ public func generate(sourceDirs: [String]?,
                         semaphore: sema,
                         queue: mockgenQueue) { (elements, imports) in
                             elements.forEach { element in
-                                protocolMap[element.name] = element
+                                protocolMap[element.entityNode.name] = element
                                 if element.isAnnotated {
-                                    annotatedProtocolMap[element.name] = element
+                                    annotatedProtocolMap[element.entityNode.name] = element
                                 }
                             }
                             if let imports = imports {
diff --git a/Sources/MockoloFramework/Operations/ProcessedTypeMapGenerator.swift b/Sources/MockoloFramework/Operations/ProcessedTypeMapGenerator.swift
index 3990fff6..57e70254 100644
--- a/Sources/MockoloFramework/Operations/ProcessedTypeMapGenerator.swift
+++ b/Sources/MockoloFramework/Operations/ProcessedTypeMapGenerator.swift
@@ -65,23 +65,11 @@ private func generateProcessedModels(_ path: String,
         let topstructure = try Structure(path: path)
         let subs = topstructure.substructures
         let results = subs.compactMap { current -> Entity? in
-            let members = current.substructures.compactMap { (child: Structure) -> Model? in
-                return Entity.model(for: child, filepath: path, data: content, overrides: nil, processed: true)
-            }
-            
-            let curAttributes = current.extractAttributes(content, filterOn: SwiftDeclarationAttributeKind.available.rawValue)
-            let hasInit = current.substructures.filter(path: \.isInitializer).count > 0
-            return Entity(name: current.name,
+            return Entity(entityNode: current,
                           filepath: path,
                           data: content,
                           isAnnotated: false,
                           overrides: nil,
-                          acl: current.accessControlLevelDescription,
-                          attributes: curAttributes,
-                          inheritedTypes: [],
-                          members: members,
-                          hasInit: hasInit,
-                          offset: current.offset,
                           isProcessed: true)
         }
         
diff --git a/Sources/MockoloFramework/Operations/ProtocolMapGenerator.swift b/Sources/MockoloFramework/Operations/ProtocolMapGenerator.swift
index 51d79bdf..c44a0c8a 100644
--- a/Sources/MockoloFramework/Operations/ProtocolMapGenerator.swift
+++ b/Sources/MockoloFramework/Operations/ProtocolMapGenerator.swift
@@ -159,30 +159,12 @@ private func generateProtcolMap(_ path: String,
             if current.isProtocol {
                 let metadata = current.annotationMetadata(with: annotationData, in: content)
                 let isAnnotated = metadata != nil
-                let members = current.substructures.compactMap { (child: Structure) -> Model? in
-                    return Entity.model(for: child, filepath: path, data: content, overrides: metadata?.typealiases, processed: false)
-                }
-                
-                var attributes = current.substructures.compactMap { (child: Structure) -> [String]? in
-                    return child.extractAttributes(content, filterOn: SwiftDeclarationAttributeKind.available.rawValue)
-                }.flatMap {$0}
-                
-                let curAttributes = current.extractAttributes(content, filterOn: SwiftDeclarationAttributeKind.available.rawValue)
-                attributes.append(contentsOf: curAttributes)
-                
-                let hasInit = current.substructures.filter(path: \.isInitializer).count > 0
                 
-                let node = Entity(name: current.name,
+                let node = Entity(entityNode: current,
                                   filepath: path,
                                   data: content,
                                   isAnnotated: isAnnotated,
                                   overrides: metadata?.typealiases,
-                                  acl: current.accessControlLevelDescription,
-                                  attributes: attributes,
-                                  inheritedTypes: current.inheritedTypes,
-                                  members: members,
-                                  hasInit: hasInit,
-                                  offset: current.offset,
                                   isProcessed: false)
                 results.append(node)
             }
diff --git a/Sources/MockoloFramework/Utils/InheritanceResolver.swift b/Sources/MockoloFramework/Utils/InheritanceResolver.swift
index 47dbaa50..4b509f48 100644
--- a/Sources/MockoloFramework/Utils/InheritanceResolver.swift
+++ b/Sources/MockoloFramework/Utils/InheritanceResolver.swift
@@ -41,17 +41,17 @@ func lookupEntities(key: String,
 
     // Look up the mock entities of a protocol specified by the name.
     if let current = protocolMap[key] {
-        
-        models.append(contentsOf: current.members)
-        if let curAttributes = current.subAttributes() {
-            attributes.append(contentsOf: curAttributes)
+        let sub = current.entityNode.subContainer(overrides: current.overrides, path: current.filepath, data: current.data, isProcessed: current.isProcessed)
+        models.append(contentsOf: sub.members)
+        if !current.isProcessed {
+            attributes.append(contentsOf: sub.attributes)
         }
         if let data = current.data {
-            pathToContents.append((current.filepath, data, current.offset))
+            pathToContents.append((current.filepath, data, current.entityNode.offset))
         }
         paths.append(current.filepath)
             // If the protocol inherits other protocols, look up their entities as well.
-            for parent in current.inheritedTypes {
+            for parent in current.entityNode.inheritedTypes {
                 if parent != .class, parent != .any, parent != .anyObject {
                     let (parentModels, parentProcessedModels, parentAttributes, parentPaths, parentPathToContents) = lookupEntities(key: parent,
                                                                                                                        protocolMap: protocolMap,
@@ -66,12 +66,13 @@ func lookupEntities(key: String,
         
     } else if let parentMock = inheritanceMap["\(key)Mock"] {
         // If the parent protocol is not in the protocol map, look it up in the input parent mocks map.
-        processedModels.append(contentsOf: parentMock.members)
-        if let parentAttributes = parentMock.subAttributes() {
-            attributes.append(contentsOf: parentAttributes)
+        let sub = parentMock.entityNode.subContainer(overrides: parentMock.overrides, path: parentMock.filepath, data: parentMock.data, isProcessed: parentMock.isProcessed)
+        processedModels.append(contentsOf: sub.members)
+        if !parentMock.isProcessed {
+            attributes.append(contentsOf: sub.attributes)
         }
         if let data = parentMock.data {
-            pathToContents.append((parentMock.filepath, data, parentMock.offset))
+            pathToContents.append((parentMock.filepath, data, parentMock.entityNode.offset))
         }
         paths.append(parentMock.filepath)
     }
diff --git a/Sources/MockoloFramework/Utils/SourceKitExtensions.swift b/Sources/MockoloFramework/Utils/SourceKitExtensions.swift
index d9c4f246..a583379b 100644
--- a/Sources/MockoloFramework/Utils/SourceKitExtensions.swift
+++ b/Sources/MockoloFramework/Utils/SourceKitExtensions.swift
@@ -24,7 +24,7 @@ struct AnnotationMetadata {
     var typealiases: [String: String]?
 }
 
-extension Structure {
+extension Structure: EntityNode {
     
     init(path: String) throws {
         self.init(sourceKitResponse: try Request.customRequest(request: [
@@ -105,6 +105,52 @@ extension Structure {
         return result
     }
     
+    func subContainer(overrides: [String: String]?, path: String?, data: Data?, isProcessed: Bool) -> EntityNodeSubContainer {
+        let memberList = members(with: path, data: data, overrides: overrides, processed: isProcessed)
+        let subAttributes = memberAttributes(with: data)
+        return EntityNodeSubContainer(attributes: subAttributes, members: memberList, hasInit: hasInitMember)
+    }
+    
+    func members(with path: String?, data: Data?, overrides: [String: String]?, processed: Bool) -> [Model] {
+        guard let path = path, let data = data else { return [] }
+        return self.substructures.compactMap { (child: Structure) -> Model? in
+            return model(for: child, filepath: path, data: data, overrides: overrides, processed: processed)
+        }
+    }
+    
+    func memberAttributes(with data: Data?) -> [String] {
+        guard let data = data else { return [] }
+        return self.substructures.compactMap { (child: Structure) -> [String]? in
+            return child.extractAttributes(data, filterOn: SwiftDeclarationAttributeKind.available.rawValue)
+        }.flatMap {$0}
+    }
+    
+    func model(for element: Structure, filepath: String, data: Data, overrides: [String: String]?, processed: Bool = false) -> Model? {
+        if element.isVariable {
+            return VariableModel(element, filepath: filepath, data: data, processed: processed)
+        } else if element.isMethod {
+            return MethodModel(element, filepath: filepath, data: data,  processed: processed)
+        } else if element.isAssociatedType {
+            return TypeAliasModel(element, filepath: filepath, data: data, overrideTypes: overrides, processed: processed)
+        }
+        
+        return nil
+    }
+
+    
+    var acl: String {
+        return accessControlLevelDescription
+    }
+    
+    var attributesDescription: String {
+        return attributes?.description ?? ""
+    }
+    
+    
+    var hasInitMember: Bool {
+        return self.substructures.filter(path: \.isInitializer).count > 0
+    }
+    
     var name: String {
         // A type must have a name.
         return dictionary["key.name"] as? String ?? .unknownVal
diff --git a/Sources/MockoloFramework/Utils/SwiftSyntaxExtensions.swift b/Sources/MockoloFramework/Utils/SwiftSyntaxExtensions.swift
index b7a99439..b14e5998 100644
--- a/Sources/MockoloFramework/Utils/SwiftSyntaxExtensions.swift
+++ b/Sources/MockoloFramework/Utils/SwiftSyntaxExtensions.swift
@@ -35,30 +35,18 @@ extension Syntax {
     var length: Int64 {
         return Int64(self.totalLength.utf8Length)
     }
-
-    func annotationMetadata(with annotation: String) -> AnnotationMetadata? {
-        return leadingTrivia?.annotationMetadata(with: annotation)
-    }
 }
 
-final class EntityVisitor: SyntaxVisitor {
-    var entities: [Entity] = []
-    var imports: [String] = []
-    let annotation: String
-    
-    init(annotation: String = "") {
-        self.annotation = annotation
-    }
-    
-    func reset() {
-        entities = []
-        imports = []
+extension AttributeListSyntax {
+    var trimmedDescription: String? {
+        return self.withoutTrivia().description.trimmingCharacters(in: .whitespacesAndNewlines)
     }
+}
 
-    // Returns access control level
-    private func acl(_  modifiers: ModifierListSyntax) -> String {
-        for m in modifiers {
-            for token in m.tokens {
+extension ModifierListSyntax {
+    var acl: String {
+        for modifier in self {
+            for token in modifier.tokens {
                 switch token.tokenKind {
                 case .publicKeyword, .internalKeyword, .privateKeyword, .fileprivateKeyword:
                     return token.text
@@ -73,78 +61,233 @@ final class EntityVisitor: SyntaxVisitor {
         }
         return ""
     }
-
-    private func attributesDescription(_ attributes: AttributeListSyntax?) -> String? {
-        return attributes?.withoutTrivia().description
+    
+    var hasStatic: Bool {
+        return self.tokens.filter {$0.tokenKind == .staticKeyword }.count > 0
     }
+}
 
-    private func hasStaticModifier(_  modifiers: ModifierListSyntax) -> Bool {
-        return modifiers.tokens.filter {$0.tokenKind == .staticKeyword }.count > 0
+extension TypeInheritanceClauseSyntax {
+    var types: [String] {
+        var list = [String]()
+        for element in self.inheritedTypeCollection {
+            if let elementName = element.firstToken?.text {
+                list.append(elementName)
+            }
+        }
+        return list
     }
     
-    private func typealiasModel(_ node: AssociatedtypeDeclSyntax, overrides: [String: String]?, acl: String, processed: Bool) -> Model {
-        // Get the inhertied type for an associated type if any
-        var t = node.inheritanceClause?.inheritedTypeCollection.description ?? ""
-        t.append(node.genericWhereClause?.description ?? "")
+    var typesDescription: String {
+        return self.inheritedTypeCollection.description
+    }
+}
+
+extension MemberDeclListSyntax {
+    
+    func memberData(with acl: String, overrides: [String: String]?, processed: Bool) -> EntityNodeSubContainer {
+        var attributeList = [String]()
+        var memberList = [Model]()
+        var hasInit = false
+        var attrDesc: String? = nil
+        for m in self {
+            if let varMember = m.decl as? VariableDeclSyntax {
+                memberList.append(contentsOf: varMember.models(with: acl, processed: processed))
+                attrDesc = varMember.attributes?.trimmedDescription
+            } else if let funcMember = m.decl as? FunctionDeclSyntax {
+                memberList.append(funcMember.model(with: acl, processed: processed))
+                attrDesc = funcMember.attributes?.trimmedDescription
+            } else if let initMember = m.decl as? InitializerDeclSyntax {
+                hasInit = true
+                memberList.append(initMember.model(with: acl, processed: processed))
+                attrDesc = initMember.attributes?.trimmedDescription
+            } else if let patMember = m.decl as? AssociatedtypeDeclSyntax {
+                memberList.append(patMember.model(with: acl, overrides: overrides, processed: processed))
+                attrDesc = patMember.attributes?.trimmedDescription
+            }
+            
+            if let attrDesc = attrDesc {
+                attributeList.append(attrDesc.trimmingCharacters(in: .whitespacesAndNewlines))
+            }
+        }
         
-        return TypeAliasModel(name: node.identifier.text,
-                              typeName: t,
-                              acl: acl,
-                              overrideTypes: overrides,
-                              offset: node.offset,
-                              length: node.length,
-                              modelDescription: node.description,
-                              processed: processed)
+        return EntityNodeSubContainer(attributes: attributeList, members: memberList, hasInit: hasInit)
+    }
+}
+
+extension ProtocolDeclSyntax: EntityNode {
+    var name: String {
+        return identifier.text
+    }
+    
+    var acl: String {
+        return self.modifiers?.acl ?? ""
+    }
+    
+    var inheritedTypes: [String] {
+        return inheritanceClause?.types ?? []
+    }
+    
+    var attributesDescription: String {
+        self.attributes?.trimmedDescription ?? ""
+    }
+    
+    var offset: Int64 {
+        return Int64(self.position.utf8Offset)
+    }
+
+    func annotationMetadata(with annotation: String) -> AnnotationMetadata? {
+        return leadingTrivia?.annotationMetadata(with: annotation)
+    }
+
+    func subContainer(overrides: [String: String]?, path: String?, data: Data?, isProcessed: Bool) -> EntityNodeSubContainer {
+        let ret = self.members.members.memberData(with: acl, overrides: overrides, processed: isProcessed)
+        return ret
+    }
+}
+
+extension ClassDeclSyntax: EntityNode {
+    
+    var name: String {
+        return identifier.text
+    }
+    
+    var acl: String {
+        return self.modifiers?.acl ?? ""
+    }
+    
+    var inheritedTypes: [String] {
+        return inheritanceClause?.types ?? []
     }
     
-    private func initModel(_ node: InitializerDeclSyntax, acl: String, processed: Bool) -> Model {
-        let params = node.parameters.parameterList.compactMap { paramModel($0, isInitializer: true) }
-        let genericTypeParams = node.genericParameterClause?.genericParameterList.compactMap { genericTypeParamModel($0, isInitializer: true) } ?? []
+    var attributesDescription: String {
+        self.attributes?.trimmedDescription ?? ""
+    }
+    
+    var offset: Int64 {
+        return Int64(self.position.utf8Offset)
+    }
+    
+    func subContainer(overrides: [String : String]?, path: String?, data: Data?, isProcessed: Bool) -> EntityNodeSubContainer {
+        let ret = self.members.members.memberData(with: acl, overrides: nil, processed: true)
+        return ret
+    }
+}
+
+extension VariableDeclSyntax {
+    func models(with acl: String, processed: Bool) -> [Model] {
+        // Detect whether it's static
+        var isStatic = false
+        if let modifiers = self.modifiers {
+            isStatic = modifiers.hasStatic
+        }
         
-        return MethodModel(name: "init",
-                           typeName: "",
-                           acl: acl,
-                           genericTypeParams: genericTypeParams,
-                           params: params,
-                           throwsOrRethrows: node.throwsOrRethrowsKeyword?.text ?? "",
-                           isStatic: false,
-                           isInitializer: true,
-                           offset: node.offset,
-                           length: node.length,
-                           modelDescription: node.description,
-                           processed: processed)
+        // Need to access pattern bindings to get name, type, and other info of a var decl
+        let varmodels = self.bindings.compactMap { (v: PatternBindingSyntax) -> Model in
+            let name = v.pattern.firstToken?.text ?? String.unknownVal
+            var typeName = ""
+            var canBeInitParam = false
+            
+            // Get the type info and whether it can be a var param for an initializer
+            if let vtype = v.typeAnnotation?.type.description {
+                typeName = vtype
+                canBeInitParam = !isStatic &&
+                    !vtype.hasSuffix("?") &&
+                    !name.hasPrefix(.underlyingVarPrefix) &&
+                    !name.hasSuffix(.closureVarSuffix) &&
+                    !name.hasSuffix(.callCountSuffix) &&
+                    !name.hasSuffix(.subjectSuffix) &&
+                    vtype != .unknownVal
+            }
+            
+            let varmodel = VariableModel(name: name,
+                                         typeName: typeName,
+                                         acl: acl,
+                                         isStatic: isStatic,
+                                         canBeInitParam: canBeInitParam,
+                                         offset: v.offset,
+                                         length: v.length,
+                                         modelDescription: self.description,
+                                         processed: processed)
+            return varmodel
+        }
+        return varmodels
     }
+}
+
+extension FunctionDeclSyntax {
     
-    private func funcModel(_ node: FunctionDeclSyntax, acl: String, processed: Bool) -> Model {
+    func model(with acl: String, processed: Bool) -> Model {
         var isStatic = false
-        if let modifiers = node.modifiers {
-            isStatic = hasStaticModifier(modifiers)
+        if let modifiers = self.modifiers {
+            isStatic = modifiers.hasStatic
         }
         
-        let params = node.signature.input.parameterList.compactMap { paramModel($0, isInitializer: false) }
-        let genericTypeParams = node.genericParameterClause?.genericParameterList.compactMap { genericTypeParamModel($0, isInitializer: false) } ?? []
+        let params = self.signature.input.parameterList.compactMap { $0.model(inInit: false) }
+        let genericTypeParams = self.genericParameterClause?.genericParameterList.compactMap { $0.model(inInit: false) } ?? []
         
-        let funcmodel = MethodModel(name: node.identifier.description,
-                                    typeName: node.signature.output?.returnType.description ?? "",
+        let funcmodel = MethodModel(name: self.identifier.description,
+                                    typeName: self.signature.output?.returnType.description ?? "",
                                     acl: acl,
                                     genericTypeParams: genericTypeParams,
                                     params: params,
-                                    throwsOrRethrows: node.signature.throwsOrRethrowsKeyword?.text ?? "",
+                                    throwsOrRethrows: self.signature.throwsOrRethrowsKeyword?.text ?? "",
                                     isStatic: isStatic,
                                     isInitializer: false,
-                                    offset: node.offset,
-                                    length: node.length,
-                                    modelDescription: node.description,
+                                    offset: self.offset,
+                                    length: self.length,
+                                    modelDescription: self.description,
                                     processed: processed)
         return funcmodel
     }
     
-    private func paramModel(_ node: FunctionParameterSyntax, isInitializer: Bool) -> ParamModel {
+    
+    
+    
+}
+
+extension InitializerDeclSyntax {
+    func model(with acl: String, processed: Bool) -> Model {
+        let params = self.parameters.parameterList.compactMap { $0.model(inInit: true) }
+        let genericTypeParams = self.genericParameterClause?.genericParameterList.compactMap { $0.model(inInit: true) } ?? []
+        
+        return MethodModel(name: "init",
+                           typeName: "",
+                           acl: acl,
+                           genericTypeParams: genericTypeParams,
+                           params: params,
+                           throwsOrRethrows: self.throwsOrRethrowsKeyword?.text ?? "",
+                           isStatic: false,
+                           isInitializer: true,
+                           offset: self.offset,
+                           length: self.length,
+                           modelDescription: self.description,
+                           processed: processed)
+    }
+    
+}
+
+
+extension GenericParameterSyntax {
+    func model(inInit: Bool) -> ParamModel {
+        return ParamModel(label: "",
+                          name: self.name.text,
+                          typeName: self.inheritedType?.description ?? "",
+                          isGeneric: true,
+                          inInit: inInit,
+                          offset: self.offset,
+                          length: self.length)
+    }
+    
+}
+
+extension FunctionParameterSyntax {
+    func model(inInit: Bool) -> ParamModel {
         var label = ""
         var name = ""
         // Get label and name of args
-        if let first = node.firstName?.text {
-            if let second = node.secondName?.text {
+        if let first = self.firstName?.text {
+            if let second = self.secondName?.text {
                 label = first
                 name = second
             } else {
@@ -158,8 +301,8 @@ final class EntityVisitor: SyntaxVisitor {
         }
         
         // Variadic args are not detected in the parser so need to manually look up
-        var type = node.type?.description ?? ""
-        if node.description.contains(type + "...") {
+        var type = self.type?.description ?? ""
+        if self.description.contains(type + "...") {
             type.append("...")
         }
         
@@ -167,87 +310,44 @@ final class EntityVisitor: SyntaxVisitor {
                           name: name,
                           typeName: type,
                           isGeneric: false,
-                          isInitializer: isInitializer,
-                          offset: node.offset,
-                          length: node.length)
+                          inInit: inInit,
+                          offset: self.offset,
+                          length: self.length)
     }
     
-    private func genericTypeParamModel(_ node: GenericParameterSyntax, isInitializer: Bool) -> ParamModel {
-        return ParamModel(label: "",
-                            name: node.name.text,
-                            typeName: node.inheritedType?.description ?? "",
-                            isGeneric: true,
-                            isInitializer: isInitializer,
-                            offset: node.offset,
-                            length: node.length)
+}
+
+extension AssociatedtypeDeclSyntax {
+    func model(with acl: String, overrides: [String: String]?, processed: Bool) -> Model {
+        // Get the inhertied type for an associated type if any
+        var t = self.inheritanceClause?.typesDescription ?? ""
+        t.append(self.genericWhereClause?.description ?? "")
+        
+        return TypeAliasModel(name: self.identifier.text,
+                              typeName: t,
+                              acl: acl,
+                              overrideTypes: overrides,
+                              offset: self.offset,
+                              length: self.length,
+                              modelDescription: self.description,
+                              processed: processed)
     }
     
-    private func varModels(_ node: VariableDeclSyntax, acl: String, processed: Bool) -> [Model] {
-        // Detect whether it's static
-        var isStatic = false
-        if let modifiers = node.modifiers {
-            isStatic = hasStaticModifier(modifiers)
-        }
-
-        // Need to access pattern bindings to get name, type, and other info of a var decl
-        let varmodels = node.bindings.compactMap { (v: PatternBindingSyntax) -> Model in
-            let name = v.pattern.firstToken?.text ?? String.unknownVal
-            var typeName = ""
-            var canBeInitParam = false
-            
-            // Get the type info and whether it can be a var param for an initializer
-            if let vtype = v.typeAnnotation?.type.description {
-                typeName = vtype
-                canBeInitParam = !isStatic &&
-                    !vtype.hasSuffix("?") &&
-                    !name.hasPrefix(.underlyingVarPrefix) &&
-                    !name.hasSuffix(.closureVarSuffix) &&
-                    !name.hasSuffix(.callCountSuffix) &&
-                    !name.hasSuffix(.subjectSuffix) &&
-                    vtype != .unknownVal
-            }
-            
-            let varmodel = VariableModel(name: name,
-                                         typeName: typeName,
-                                         acl: acl,
-                                         isStatic: isStatic,
-                                         canBeInitParam: canBeInitParam,
-                                         offset: v.offset,
-                                         length: v.length,
-                                         modelDescription: node.description,
-                                         processed: processed)
-            return varmodel
-        }
-        return varmodels
+    
+    
+}
+final class EntityVisitor: SyntaxVisitor {
+    var entities: [Entity] = []
+    var imports: [String] = []
+    let annotation: String
+    
+    init(annotation: String = "") {
+        self.annotation = annotation
     }
     
-    private func memberList(_ members: MemberDeclListSyntax, overrides: [String: String]?, acl: String, processed: Bool) -> ([String], [Model], Bool) {
-        var attributeList = [String]()
-        var memberList = [Model]()
-        var hasInit = false
-        var attrDesc: String? = nil
-        for m in members {
-            if let varMember = m.decl as? VariableDeclSyntax {
-                memberList.append(contentsOf: varModels(varMember, acl: acl, processed: processed))
-                attrDesc = attributesDescription(varMember.attributes)
-            } else if let funcMember = m.decl as? FunctionDeclSyntax {
-                memberList.append(funcModel(funcMember, acl: acl, processed: processed))
-                attrDesc = attributesDescription(funcMember.attributes)
-            } else if let initMember = m.decl as? InitializerDeclSyntax {
-                hasInit = true
-                memberList.append(initModel(initMember, acl: acl, processed: processed))
-                attrDesc = attributesDescription(initMember.attributes)
-            } else if let patMember = m.decl as? AssociatedtypeDeclSyntax {
-                memberList.append(typealiasModel(patMember, overrides: overrides, acl: acl, processed: processed))
-                attrDesc = attributesDescription(patMember.attributes)
-            }
-            
-            if let attrDesc = attrDesc {
-                attributeList.append(attrDesc.trimmingCharacters(in: .whitespacesAndNewlines))
-            }
-        }
-        
-        return (attributeList, memberList, hasInit)
+    func reset() {
+        entities = []
+        imports = []
     }
     
     func visit(_ node: ProtocolDeclSyntax) -> SyntaxVisitorContinueKind {
@@ -259,64 +359,16 @@ final class EntityVisitor: SyntaxVisitor {
             overrides = metadata?.typealiases
         }
         
-        var parentList = [String]()
-        if let parents = node.inheritanceClause?.inheritedTypeCollection {
-            for p in parents {
-                if let pname = p.firstToken?.text {
-                    parentList.append(pname)
-                }
-            }
-        }
-        
-        var aclDesc = ""
-        if let mds = node.modifiers {
-            aclDesc = acl(mds)
-        }
-        
-        let (attributes, members, hasInit) = memberList(node.members.members, overrides: overrides, acl: aclDesc, processed: false)
-        
-        var attributeList = attributes
-        if let attrDesc = node.attributes?.withoutTrivia().description {
-            attributeList.append(attrDesc.trimmingCharacters(in: .whitespacesAndNewlines))
-        }
-        
-        let ent = Entity(name: node.identifier.text,
+        let ent = Entity(entityNode: node,
                          isAnnotated: isAnnotated,
                          overrides: overrides,
-                         acl: aclDesc,
-                         attributes: attributeList,
-                         inheritedTypes: parentList,
-                         members: members,
-                         hasInit: hasInit,
-                         offset: node.offset,
                          isProcessed: false)
         entities.append(ent)
         return .skipChildren
     }
     
     func visit(_ node: ClassDeclSyntax) -> SyntaxVisitorContinueKind {
-        var aclDesc = ""
-        if let mds = node.modifiers {
-            aclDesc = acl(mds)
-        }
-        
-        let (attributes, members, hasInit) = memberList(node.members.members, overrides: nil, acl: aclDesc, processed: true)
-        
-        var attributeList = attributes
-        if let attrDesc = node.attributes?.withoutTrivia().description {
-            attributeList.append(attrDesc.trimmingCharacters(in: .whitespacesAndNewlines))
-        }
-        
-        let ent = Entity(name: node.identifier.text,
-                         isAnnotated: false,
-                         overrides: nil,
-                         acl: aclDesc,
-                         attributes: attributeList,
-                         inheritedTypes: [],
-                         members: members,
-                         hasInit: hasInit,
-                         offset: node.offset,
-                         isProcessed: false)
+        let ent = Entity(entityNode: node, isAnnotated: false, overrides: nil, isProcessed: true)
         entities.append(ent)
         return .skipChildren
     }
diff --git a/Tests/MockoloTestCase.swift b/Tests/MockoloTestCase.swift
index 37524546..70337c39 100644
--- a/Tests/MockoloTestCase.swift
+++ b/Tests/MockoloTestCase.swift
@@ -106,6 +106,7 @@ class MockoloTestCase: XCTestCase {
         \(macroEnd)
         """
         
+        let useSourceKit = Int.random(in: 0..<10) > 5
         try? generate(sourceDirs: nil,
                       sourceFiles: srcFilePaths,
                       exclusionSuffixes: ["Mocks", "Tests"],
@@ -113,7 +114,7 @@ class MockoloTestCase: XCTestCase {
                       annotation: String.mockAnnotation,
                       header: header,
                       macro: "MOCK",
-                      parserType: .sourceKit,
+                      parserType: useSourceKit ? .sourceKit : .swiftSyntax,
                       to: dstFilePath,
                       loggingLevel: 3,
                       concurrencyLimit: concurrencyLimit,