Unify #if handling and fix top-level #if regression

Sources/MockoloFramework/Models/ConditionalBlock.swift

@@ -0,0 +1,49 @@
+//
+//  Copyright (c) 2018. Uber Technologies
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+/// One generic conditional-compilation structure (`#if`/`#elseif`/`#else`), specialized by leaf
+/// type. Member-level (`IfMacroModel`), import-level, and top-level `#if` handling all reuse this
+/// so the clause/branch shape — and directive emission — live in exactly one place.
+struct ConditionalBlock<Leaf> {
+    /// A single `#if` / `#elseif` / `#else` clause and the leaves declared under it.
+    struct Clause {
+        var type: IfClauseType
+        var elements: [Leaf]
+    }
+
+    let clauses: [Clause]
+    let offset: Int64
+}
+
+extension IfClauseType {
+    /// The directive line (`#if`/`#elseif`/`#else`) for this clause at the given indent level.
+    /// Centralizes what used to be duplicated switches in `IfMacroTemplate` and `ImportsHandler`.
+    func directiveLine(indent: Int) -> String {
+        switch self {
+        case .if(let condition):
+            return "\(indent.tab)#if \(condition)"
+        case .elseif(let condition):
+            return "\(indent.tab)#elseif \(condition)"
+        case .else:
+            return "\(indent.tab)#else"
+        }
+    }
+
+    /// The closing `#endif` line at the given indent level.
+    static func endIfLine(indent: Int) -> String {
+        "\(indent.tab)#endif"
+    }
+}

Sources/MockoloFramework/Models/ConditionalImportBlock.swift

@@ -14,25 +14,9 @@
 //  limitations under the License.
 //
 
-/// Represents import content: either a simple import statement or a nested conditional block
+/// Represents import content: either a simple import statement or a nested conditional block.
+/// Conditional import blocks reuse the shared `ConditionalBlock` structure (leaf == `ImportContent`).
 indirect enum ImportContent {
     case simple(Import)
-    case conditional(ConditionalImportBlock)
-}
-
-/// Represents a conditional import block (#if/#elseif/#else/#endif)
-struct ConditionalImportBlock {
-    /// Represents a single clause in a conditional import block
-    struct Clause {
-        var type: IfClauseType
-        var contents: [ImportContent]
-    }
-
-    let clauses: [Clause]
-    let offset: Int64
-
-    init(clauses: [Clause], offset: Int64) {
-        self.clauses = clauses
-        self.offset = offset
-    }
+    case conditional(ConditionalBlock<ImportContent>)
 }

Sources/MockoloFramework/Models/IfMacroModel.swift

@@ -31,39 +31,40 @@ enum IfClauseType {
 }
 
 final class IfMacroModel: Model {
-    /// Represents a single clause in a conditional compilation block
-    struct Clause {
-        var type: IfClauseType
-        var entities: [(String, Model)]
-    }
+    /// A member-level conditional-compilation clause: `(overloadingResolvedName, Model)` leaves.
+    typealias Clause = ConditionalBlock<(String, Model)>.Clause
+
+    let block: ConditionalBlock<(String, Model)>
 
-    let clauses: [Clause]
-    let offset: Int64
+    var clauses: [Clause] { block.clauses }
+    var offset: Int64 { block.offset }
 
     var modelType: ModelType {
         .macro
     }
-    
+
     var name: String {
-        clauses.first?.type.condition ?? ""
+        block.clauses.first?.type.condition ?? ""
     }
 
     var fullName: String {
-        clauses.flatMap(\.entities).map { $0.0 }.joined(separator: "_")
+        block.clauses.flatMap(\.elements).map { $0.0 }.joined(separator: "_")
+    }
+
+    init(block: ConditionalBlock<(String, Model)>) {
+        self.block = block
     }
 
     /// Creates an IfMacroModel with multiple clauses
-    init(clauses: [Clause], offset: Int64) {
-        self.clauses = clauses
-        self.offset = offset
+    convenience init(clauses: [Clause], offset: Int64) {
+        self.init(block: ConditionalBlock(clauses: clauses, offset: offset))
     }
 
     /// Initializer for simple #if blocks
     convenience init(name: String,
                      offset: Int64,
                      entities: [(String, Model)]) {
-        let clause = Clause(type: .if(name), entities: entities)
-        self.init(clauses: [clause], offset: offset)
+        self.init(clauses: [Clause(type: .if(name), elements: entities)], offset: offset)
     }
 
     func render(

Sources/MockoloFramework/Models/TopLevelConditionalLeaf.swift

@@ -0,0 +1,85 @@
+//
+//  Copyright (c) 2018. Uber Technologies
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+/// A leaf of a top-level `#if` block. Unlike member-level (`Model`) or import-only
+/// (`ImportContent`) blocks, a top-level clause can mix imports and declarations, so its leaf is a
+/// small union that preserves source order within the clause.
+indirect enum TopLevelConditionalLeaf {
+    case `import`(Import)
+    case entity(Entity)
+    case nested(ConditionalBlock<TopLevelConditionalLeaf>)
+}
+
+extension ConditionalBlock where Leaf == TopLevelConditionalLeaf {
+    /// Whether any clause (recursively) declares a mockable entity.
+    var containsEntities: Bool {
+        clauses.contains { clause in
+            clause.elements.contains { leaf in
+                switch leaf {
+                case .entity:
+                    return true
+                case .nested(let nested):
+                    return nested.containsEntities
+                case .import:
+                    return false
+                }
+            }
+        }
+    }
+
+    /// All entity offsets declared in this block, so flat rendering can skip the entities that are
+    /// rendered (wrapped in their `#if` directives) by `renderConditionalBlocks`.
+    func entityOffsets() -> [Int64] {
+        clauses.flatMap { clause in
+            clause.elements.flatMap { leaf -> [Int64] in
+                switch leaf {
+                case .entity(let entity):
+                    return [entity.entityNode.offset]
+                case .nested(let nested):
+                    return nested.entityOffsets()
+                case .import:
+                    return []
+                }
+            }
+        }
+    }
+
+    /// Lowers to the import pipeline, preserving the `#if`/`#elseif`/`#else` structure so
+    /// conditional imports keep their directives. Returns nil when no clause carries an import, so a
+    /// declaration-only block never emits a stray empty import block.
+    func loweredToImports() -> ConditionalBlock<ImportContent>? {
+        var hasImport = false
+        let importClauses = clauses.map { clause -> ConditionalBlock<ImportContent>.Clause in
+            var contents: [ImportContent] = []
+            for leaf in clause.elements {
+                switch leaf {
+                case .import(let imp):
+                    contents.append(.simple(imp))
+                    hasImport = true
+                case .nested(let nested):
+                    if let lowered = nested.loweredToImports() {
+                        contents.append(.conditional(lowered))
+                        hasImport = true
+                    }
+                case .entity:
+                    break
+                }
+            }
+            return ConditionalBlock<ImportContent>.Clause(type: clause.type, elements: contents)
+        }
+        return hasImport ? ConditionalBlock<ImportContent>(clauses: importClauses, offset: offset) : nil
+    }
+}

Sources/MockoloFramework/Operations/Generator.swift

@@ -58,12 +58,13 @@ public func generate(sourceDirs: [String],
     var annotatedProtocolMap = [String: Entity]()
     var pathToImportsMap = ImportMap()
     var relevantPaths = [String]()
+    var topLevelConditionalBlocks = [ConditionalBlock<TopLevelConditionalLeaf>]()
 
     signpost_begin(name: "Process input")
     let t0 = CFAbsoluteTimeGetCurrent()
     log("Process input mock files...", level: .info)
     if let mockFilePaths = mockFilePaths, !mockFilePaths.isEmpty {
-        parser.parseProcessedDecls(mockFilePaths, fileMacro: macro) { (elements, imports) in
+        parser.parseProcessedDecls(mockFilePaths, fileMacro: macro) { (elements, imports, _) in
                                     elements.forEach { element in
                                         parentMocks[element.entityNode.nameText] = element
                                     }
@@ -94,7 +95,7 @@ public func generate(sourceDirs: [String],
                       exclusionSuffixes: exclusionSuffixes,
                       annotation: annotation,
                       fileMacro: macro,
-                      declType: declType) { (elements, imports) in
+                      declType: declType) { (elements, imports, conditionals) in
                         elements.forEach { element in
                             protocolMap[element.entityNode.nameText] = element
                             if element.isAnnotated {
@@ -106,6 +107,7 @@ public func generate(sourceDirs: [String],
                                 pathToImportsMap[path] = importMap
                             }
                         }
+                        topLevelConditionalBlocks.append(contentsOf: conditionals)
     }
     signpost_end(name: "Generate protocol map")
     let t2 = CFAbsoluteTimeGetCurrent()
@@ -148,15 +150,31 @@ public func generate(sourceDirs: [String],
 
     signpost_begin(name: "Render models")
     log("Render models with templates...", level: .info)
+    let arguments = GenerationArguments(
+        useTemplateFunc: useTemplateFunc,
+        allowSetCallCount: allowSetCallCount,
+        mockFinal: mockFinal,
+        enableFuncArgsHistory: enableFuncArgsHistory,
+        disableCombineDefaultValues: disableCombineDefaultValues
+    )
+    // Entities declared inside a top-level `#if` are rendered (wrapped in their directives) by
+    // `renderConditionalBlocks`; skip them in the flat pass so each mock is emitted exactly once.
+    let conditionalOffsets = Set(topLevelConditionalBlocks.flatMap { $0.entityOffsets() })
+    let resolvedByOffset = Dictionary(
+        resolvedEntities.map { ($0.entity.entityNode.offset, $0) },
+        uniquingKeysWith: { first, _ in first }
+    )
+    let flatEntities = resolvedEntities.filter { !conditionalOffsets.contains($0.entity.entityNode.offset) }
     renderTemplates(
-        entities: resolvedEntities,
-        arguments: .init(
-            useTemplateFunc: useTemplateFunc,
-            allowSetCallCount: allowSetCallCount,
-            mockFinal: mockFinal,
-            enableFuncArgsHistory: enableFuncArgsHistory,
-            disableCombineDefaultValues: disableCombineDefaultValues
-        )
+        entities: flatEntities,
+        arguments: arguments
+    ) { (mockString: String, offset: Int64) in
+                        candidates.append((mockString, offset))
+    }
+    renderConditionalBlocks(
+        topLevelConditionalBlocks,
+        resolvedByOffset: resolvedByOffset,
+        arguments: arguments
     ) { (mockString: String, offset: Int64) in
                         candidates.append((mockString, offset))
     }

Sources/MockoloFramework/Operations/ImportsHandler.swift

@@ -22,7 +22,7 @@ func handleImports(pathToImportsMap: ImportMap,
                    testableImports: [String]?,
                    relevantPaths: [String]) -> String {
     var topLevelImports: [Import] = []
-    var conditionalBlocks: [ConditionalImportBlock] = []
+    var conditionalBlocks: [ConditionalBlock<ImportContent>] = []
 
     // 1. Collect imports from all relevant files
     for (path, parsedImports) in pathToImportsMap {
@@ -99,19 +99,12 @@ private func renderImportContents(
 
             var result = ""
             for clause in block.clauses {
-                switch clause.type {
-                case .if(let condition):
-                    result += "#if \(condition)\n"
-                case .elseif(let condition):
-                    result += "#elseif \(condition)\n"
-                case .else:
-                    result += "#else\n"
-                }
+                result += clause.type.directiveLine(indent: 0) + "\n"
                 // Recursively render nested block
-                result += renderImportContents(clause.contents, excludeImports: excludeImports, testableImports: testableImports)
+                result += renderImportContents(clause.elements, excludeImports: excludeImports, testableImports: testableImports)
                 result += "\n"
             }
-            result += "#endif"
+            result += IfClauseType.endIfLine(indent: 0)
             clauseLines.append(result)
         }
     }
@@ -126,7 +119,7 @@ private func visitModuleName(_ contents: [ImportContent]) -> [String] {
         case .simple(let `import`):
             return [`import`.moduleName]
         case .conditional(let block):
-            return visitModuleName(block.clauses.flatMap(\.contents))
+            return visitModuleName(block.clauses.flatMap(\.elements))
         }
     }
 }

Sources/MockoloFramework/Operations/TemplateRenderer.swift

@@ -31,3 +31,58 @@ func renderTemplates(entities: [ResolvedEntity],
         }
     }
 }
+
+/// Renders mocks for entities declared inside top-level `#if` blocks, re-wrapping each in the
+/// original `#if`/`#elseif`/`#else` directives so the generated file compiles under the same
+/// build configurations as the source.
+func renderConditionalBlocks(_ blocks: [ConditionalBlock<TopLevelConditionalLeaf>],
+                             resolvedByOffset: [Int64: ResolvedEntity],
+                             arguments: GenerationArguments,
+                             completion: (String, Int64) -> ()) {
+    for block in blocks {
+        if let rendered = renderConditionalBlock(block,
+                                                 resolvedByOffset: resolvedByOffset,
+                                                 arguments: arguments,
+                                                 indent: 0) {
+            completion(rendered, block.offset)
+        }
+    }
+}
+
+private func renderConditionalBlock(_ block: ConditionalBlock<TopLevelConditionalLeaf>,
+                                    resolvedByOffset: [Int64: ResolvedEntity],
+                                    arguments: GenerationArguments,
+                                    indent: Int) -> String? {
+    var lines: [String] = []
+    var produced = false
+    for clause in block.clauses {
+        // Keep every clause so #elseif/#else stay syntactically valid even when a branch is empty.
+        lines.append(clause.type.directiveLine(indent: indent))
+        for leaf in clause.elements {
+            switch leaf {
+            case .import:
+                // Imports are emitted by the imports section (ImportsHandler), not here.
+                break
+            case .entity(let entity):
+                if let resolved = resolvedByOffset[entity.entityNode.offset],
+                   let mock = resolved.model().render(context: .init(), arguments: arguments),
+                   !mock.isEmpty {
+                    lines.append(mock)
+                    produced = true
+                }
+            case .nested(let nested):
+                if let rendered = renderConditionalBlock(nested,
+                                                         resolvedByOffset: resolvedByOffset,
+                                                         arguments: arguments,
+                                                         indent: indent) {
+                    lines.append(rendered)
+                    produced = true
+                }
+            }
+        }
+    }
+    // Drop a block that produced no mock (e.g. only non-annotated decls) rather than emit an empty #if.
+    guard produced else { return nil }
+    lines.append(IfClauseType.endIfLine(indent: indent))
+    return lines.joined(separator: "\n")
+}