1 files changed, 37 insertions, 0 deletions
diff --git a/Sources/NorgEditor/NorgAutoIndent.swift b/Sources/NorgEditor/NorgAutoIndent.swift
new file mode 100644
index 0000000..2e6fa4b
--- /dev/null
+++ b/Sources/NorgEditor/NorgAutoIndent.swift
@@ -0,0 +1,37 @@
+import Foundation
+import NorgKit
+
+/// Computes the leading whitespace to insert when the operator presses return.
+public enum NorgAutoIndent {
+
+  /// The whitespace to insert after a newline.
+  public static func indentation(for text: String, newlineAt location: Int) -> String {
+    let nsText = text as NSString
+    guard location >= 0, location <= nsText.length else { return "" }
+
+    let lineRange = nsText.lineRange(for: NSRange(location: location, length: 0))
+    let prefix = nsText.substring(to: NSMaxRange(lineRange))
+    return indentation(forPrefix: prefix)
+  }
+
+  static func indentation(forPrefix prefix: String) -> String {
+    var openHeadingLevels: [Int] = []
+    for token in NorgLexer.tokenize(prefix) {
+      switch token.kind {
+      case .heading(let level):
+        while let deepest = openHeadingLevels.last, deepest >= level {
+          openHeadingLevels.removeLast()
+        }
+        openHeadingLevels.append(level)
+      case .strongDelimiter:
+        openHeadingLevels.removeAll()
+      case .weakDelimiter:
+        if !openHeadingLevels.isEmpty { openHeadingLevels.removeLast() }
+      default:
+        break
+      }
+    }
+    guard let level = openHeadingLevels.last else { return "" }
+    return String(repeating: " ", count: level + 1)
+  }
+}