Skip to content

Commit 8d81f28

Browse files
davdromandavdroman
committed
WIP
1 parent c0db8af commit 8d81f28

File tree

5 files changed

+675
-43
lines changed

5 files changed

+675
-43
lines changed

‎Sources/Once/Exports.swift

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
@_exported public import Foundation

‎Sources/Once/Macros.swift

+6
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
/// A macro that ensures the specified block is only executed once per program run.
12
@discardableResult
23
@freestanding(expression)
34
public macro once<T>(block: () throws -> T) -> T? = #externalMacro(module: "OnceMacro", type: "OnceMacro")
5+
6+
/// A macro that ensures the specified block is only executed once per program run.
7+
@discardableResult
8+
@freestanding(expression)
9+
public macro once<T>(block: () async throws -> T) -> T? = #externalMacro(module: "OnceMacro", type: "OnceMacro")

‎Sources/OnceMacro/OnceMacro.swift

+110-15
Original file line numberDiff line numberDiff line change
@@ -14,23 +14,118 @@ public struct OnceMacro: ExpressionMacro {
1414
tryFinder.walk(block)
1515
let tryKeyword = tryFinder.found ? "try " : ""
1616

17-
return """
18-
{
19-
@MainActor
20-
enum Once {
21-
private static var hasRun: Bool = false
22-
23-
@discardableResult
24-
static func run<T>(_ block: () throws -> T) rethrows -> T? {
25-
guard !hasRun else { return nil }
26-
hasRun = true
27-
return try block()
17+
let awaitFinder = AwaitExprFinder()
18+
awaitFinder.walk(block)
19+
// let awaitKeyword = awaitFinder.found ? "await " : ""
20+
21+
if !awaitFinder.found {
22+
return """
23+
{
24+
final class Once: @unchecked Sendable {
25+
static let shared = Once()
26+
27+
private var hasRun = false
28+
private var isRunning = false
29+
private let lock = NSLock()
30+
31+
@discardableResult
32+
func run<T>(_ block: () throws -> T) rethrows -> T? {
33+
lock.lock()
34+
if hasRun || isRunning {
35+
lock.unlock()
36+
return nil
37+
}
38+
isRunning = true
39+
lock.unlock()
40+
41+
defer {
42+
lock.lock()
43+
isRunning = false
44+
hasRun = true
45+
lock.unlock()
46+
}
47+
48+
return try block()
49+
}
2850
}
29-
}
3051
31-
return \(raw: tryKeyword)Once.run \(block)
32-
}()
33-
"""
52+
return \(raw: tryKeyword)Once.shared.run \(block)
53+
}()
54+
"""
55+
} else {
56+
return """
57+
{
58+
actor Once {
59+
static let shared = Once()
60+
61+
private var hasRun = false
62+
private var isRunning = false
63+
64+
@discardableResult
65+
func run<T: Sendable>(_ block: () async throws -> T) async rethrows -> T? {
66+
guard !hasRun && !isRunning else {
67+
return nil
68+
}
69+
isRunning = true
70+
71+
defer {
72+
isRunning = false
73+
hasRun = true
74+
}
75+
76+
return try await block()
77+
}
78+
}
79+
80+
return \(raw: tryKeyword)await Once.shared.run \(block)
81+
}()
82+
"""
83+
}
84+
85+
// return """
86+
// {
87+
// @MainActor
88+
// enum Once {
89+
// private static var hasRun: Bool = false
90+
//
91+
// @discardableResult
92+
// static func run<T>(_ block: () throws -> T) rethrows -> T? {
93+
// guard !hasRun else { return nil }
94+
// hasRun = true
95+
// return try block()
96+
// }
97+
// }
98+
//
99+
// return \(raw: tryKeyword)Once.run \(block)
100+
// }()
101+
// """
102+
}
103+
}
104+
105+
final class AwaitExprFinder: SyntaxVisitor {
106+
private var atRoot = true
107+
private(set) var found = false
108+
109+
init() {
110+
super.init(viewMode: .sourceAccurate)
111+
}
112+
113+
override func visit(_ node: AwaitExprSyntax) -> SyntaxVisitorContinueKind {
114+
found = true
115+
return .skipChildren
116+
}
117+
118+
override func visit(_ node: FunctionDeclSyntax) -> SyntaxVisitorContinueKind {
119+
.skipChildren
120+
}
121+
122+
override func visit(_ node: ClosureExprSyntax) -> SyntaxVisitorContinueKind {
123+
if atRoot {
124+
atRoot = false
125+
return .visitChildren
126+
} else {
127+
return .skipChildren
128+
}
34129
}
35130
}
36131

‎Tests/OnceMacroTests/OnceMacroTests.swift

+63-21
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,35 @@ struct OnceMacroTests {
2121
} expansion: {
2222
"""
2323
{
24-
@MainActor
25-
enum Once {
26-
private static var hasRun: Bool = false
24+
final class Once: @unchecked Sendable {
25+
static let shared = Once()
26+
27+
private var hasRun = false
28+
private var isRunning = false
29+
private let lock = NSLock()
2730
2831
@discardableResult
29-
static func run<T>(_ block: () throws -> T) rethrows -> T? {
30-
guard !hasRun else {
32+
func run<T>(_ block: () throws -> T) rethrows -> T? {
33+
lock.lock()
34+
if hasRun || isRunning {
35+
lock.unlock()
3136
return nil
3237
}
33-
hasRun = true
38+
isRunning = true
39+
lock.unlock()
40+
41+
defer {
42+
lock.lock()
43+
isRunning = false
44+
hasRun = true
45+
lock.unlock()
46+
}
47+
3448
return try block()
3549
}
3650
}
3751
38-
return Once.run {
52+
return Once.shared.run {
3953
print("Hello, world!")
4054
}
4155
}()
@@ -54,21 +68,35 @@ struct OnceMacroTests {
5468
} expansion: {
5569
"""
5670
try {
57-
@MainActor
58-
enum Once {
59-
private static var hasRun: Bool = false
71+
final class Once: @unchecked Sendable {
72+
static let shared = Once()
73+
74+
private var hasRun = false
75+
private var isRunning = false
76+
private let lock = NSLock()
6077
6178
@discardableResult
62-
static func run<T>(_ block: () throws -> T) rethrows -> T? {
63-
guard !hasRun else {
79+
func run<T>(_ block: () throws -> T) rethrows -> T? {
80+
lock.lock()
81+
if hasRun || isRunning {
82+
lock.unlock()
6483
return nil
6584
}
66-
hasRun = true
85+
isRunning = true
86+
lock.unlock()
87+
88+
defer {
89+
lock.lock()
90+
isRunning = false
91+
hasRun = true
92+
lock.unlock()
93+
}
94+
6795
return try block()
6896
}
6997
}
7098
71-
return try Once.run {
99+
return try Once.shared.run {
72100
try print("Hello, world!")
73101
}
74102
}()
@@ -100,21 +128,35 @@ struct OnceMacroTests {
100128
} expansion: {
101129
"""
102130
{
103-
@MainActor
104-
enum Once {
105-
private static var hasRun: Bool = false
131+
final class Once: @unchecked Sendable {
132+
static let shared = Once()
133+
134+
private var hasRun = false
135+
private var isRunning = false
136+
private let lock = NSLock()
106137
107138
@discardableResult
108-
static func run<T>(_ block: () throws -> T) rethrows -> T? {
109-
guard !hasRun else {
139+
func run<T>(_ block: () throws -> T) rethrows -> T? {
140+
lock.lock()
141+
if hasRun || isRunning {
142+
lock.unlock()
110143
return nil
111144
}
112-
hasRun = true
145+
isRunning = true
146+
lock.unlock()
147+
148+
defer {
149+
lock.lock()
150+
isRunning = false
151+
hasRun = true
152+
lock.unlock()
153+
}
154+
113155
return try block()
114156
}
115157
}
116158
117-
return Once.run {
159+
return Once.shared.run {
118160
try! doSomething() {
119161
try foo()
120162
}

0 commit comments

Comments
 (0)