swift-pterodactyl / Sources / PterodactylSyntax / Parser.swift

// SPDX-FileCopyrightText: 2025 The Project Pterodactyl Developers
// SPDX-FileCopyrightText: 2026 The Project Pterodactyl Developers
//
// SPDX-License-Identifier: MPL-2.0

import OrderedCollections
import PterodactylFeedback

enum PrecedenceGroup {
	case equality
	case arrow

	enum Associativity {
		case right
		case none
	}

	var associativity: Associativity {
		switch self {
		case .equality: return .none
		case .arrow: return .right
		}
	}

	enum Relation {
		case tighter
		case looser
		case same
		case unrelated
	}

	private var ordinal: Int {
		switch self {
		case .arrow: return 0
		case .equality: return 1
		}
	}

	func relation(to other: PrecedenceGroup) -> Relation {
		if self == other { return .same }
		return ordinal > other.ordinal ? .tighter : .looser
	}

	func canBind(above lowerBound: PrecedenceGroup?) -> Bool {
		guard let lowerBound else { return true }
		switch relation(to: lowerBound) {
		case .tighter: return true
		case .looser: return false
		case .same: return lowerBound.associativity == .right
		case .unrelated: return true
		}
	}
}


struct Scanner {
	let tokens: [Token]
	var position: Int

	var currentToken: Token? {
		if tokens.indices.contains(position) {
			tokens[position]
		} else {
			nil
		}
	}

	mutating func eat(kindSatisfying predicate: (TokenKind) -> Bool) -> Bool {
		guard let currentToken, predicate(currentToken.kind) else { return false }
		position += 1
		return true
	}

	mutating func eat(kind: TokenKind) -> Bool {
		eat(kindSatisfying: { $0 == kind })
	}

	mutating func eatTrivia() {
		while let currentToken, currentToken.kind.isTrivia { position += 1 }
	}

	public func isAt(kind: TokenKind) -> Bool {
		currentToken?.kind == kind
	}

	public func isAt(kindSatisfying predicate: (TokenKind) -> Bool) -> Bool {
		guard let currentToken else { return false }
		return predicate(currentToken.kind)
	}
}


extension Parser {
	var scanner: Scanner {
		Scanner(tokens: tokens, position: position)
	}
}


struct DelimiterInfo {
	let left: TokenKind
	let sep: TokenKind
	let right: TokenKind

	static var block: DelimiterInfo {
		DelimiterInfo(left: .blockBegin, sep: .blockSep, right: .blockEnd)
	}

	static var curly: DelimiterInfo {
		DelimiterInfo(left: .punctuation(.lbrace), sep: .punctuation(.semicolon), right: .punctuation(.rbrace))
	}

	var leftProduction: Production {
		Parser.token(kind: left)
			.then { state in
				_ = separatorProduction.tryParse(&state)
			}
	}

	var separatorProduction: Production {
		Parser.token(kind: sep).precededByTrivia.many1
	}

	var rightProduction: Production {
		Parser.token(kind: right).precededByTrivia
	}
}

extension Parser {
	mutating func keyword(_ keyword: Keyword) {
		expect(kind: .keyword(keyword), metadata: TokenMetadata(semanticTokenType: .keyword))
	}
}

struct Production {
	var name: String
	var _precondition: (_ state: Scanner) -> Bool
	var _commit: (_ state: inout Parser) -> Void

	init(name: String, precondition: @escaping (_: Scanner) -> Bool, commit: @escaping (_: inout Parser) -> Void) {
		self.name = name
		self._precondition = precondition
		self._commit = commit
	}

	init(
		name: String,
		kind: SyntaxTreeKind,
		metadata: SyntaxTreeMetadata? = nil,
		precondition: @escaping (_: Scanner) -> Bool,
		commit: @escaping (_: inout Parser) -> Void
	) {
		self.init(name: name) { state in
			precondition(state)
		} commit: { state in
			state.open()
			defer { state.close(kind: kind, metadata: metadata) }
			commit(&state)
		}
	}

	func then(_ body: @escaping (_ state: inout Parser) -> Void) -> Production {
		Production(name: "\(name).then") { scanner in
			self._precondition(scanner)
		} commit: { state in
			self._commit(&state)
			body(&state)
		}
	}

	func precondition(_ state: Scanner) -> Bool {
		_precondition(state)
	}

	func commit(_ state: inout Parser) {
		let position = state.position
		_commit(&state)
		assert(state.position > position, "Parser `\(name)` did not consume any tokens.")
	}

	func tryParse(_ state: inout Parser, error: String? = nil) -> Bool {
		guard precondition(state.scanner) else {
			if let error {
				state.diagnostics.append(Diagnostic(message: error, severity: .error, utf16Range: state.absoluteRangeOfCurrentToken))
			}
			return false
		}
		commit(&state)
		return true
	}

	func parse(_ state: inout Parser, error: String) -> Bool {
		if !tryParse(&state) {
			state.recover(until: state.recoveryStack.current, expecting: [], message: error)
			return false
		}

		return true
	}

	var precededByTrivia: Production {
		Production(name: self.name + ".precededByTrivia") { scanner in
			var scanner = scanner
			scanner.eatTrivia()
			return precondition(scanner)
		} commit: { state in
			state.eatTrivia()
			_commit(&state)
		}
	}

	func wrap(name: String, kind: SyntaxTreeKind, metadata: SyntaxTreeMetadata? = nil) -> Production {
		Production(name: name, kind: kind, metadata: metadata, precondition: _precondition, commit: _commit)
	}

	static func cascade(_ productions: [Production]) -> Production {
		Production(name: "cascade[_\(productions.map(\.name))]") { scanner in
			productions.contains { $0.precondition(scanner) }
		} commit: { state in
			for production in productions {
				if production.tryParse(&state) {
					return
				}
			}
		}
	}

	func delimitedList(_ delim: DelimiterInfo) -> Production {
		delim.leftProduction.then { state in
			state.recoveryStack.push([.blockSep, .blockEnd])
			defer { state.recoveryStack.pop() }
			defer { _ = delim.rightProduction.parse(&state, error: "Missing closing delimiter: \(delim.right)") }
			while true {
				_ = precededByTrivia.tryParse(&state)
				guard delim.separatorProduction.tryParse(&state) else { break }
			}
		}
	}

	var many1: Production {
		return then { state in
			while precededByTrivia.tryParse(&state) {}
		}
	}
}

extension Parser {
	static func token(kind: TokenKind, metadata: TokenMetadata? = nil) -> Production {
		Production(name: "token[\(kind)]") { scanner in
			scanner.isAt(kind: kind)
		} commit: { state in
			state.expect(kind: kind, metadata: metadata)
		}
	}

	private mutating func typedBinder() {
		recoveryStack.push([.punctuation(.lparen), .punctuation(.colon), .punctuation(.rparen)])
		defer { recoveryStack.pop() }

		expect(kind: .punctuation(.lparen))
		let gotIdentifier = Self.token(kind: .identifier).precededByTrivia.tryParse(&self, error: "Binder: expected identifier here")
		let gotColon = Self.token(kind: .punctuation(.colon)).precededByTrivia.tryParse(&self, error: gotIdentifier ? "Binder: expected type annotation here" : nil)
		let gotTypeExpr = gotColon ? Self.expression.precededByTrivia.parse(&self, error: "Binder: expected type expression here") : false
		_ = Self.token(kind: .punctuation(.rparen)).precededByTrivia.tryParse(&self, error: gotTypeExpr ? "Binder: expected closing parenthesis" : nil)
	}

	static var funTypeBinder: Production {
		Production(name: "funTypeBinder", kind: .typedBinder) { scanner in
			var scanner = scanner
			guard scanner.eat(kind: .punctuation(.lparen)) else { return false }
			scanner.eatTrivia()
			guard scanner.eat(kind: .identifier) else { return false }
			scanner.eatTrivia()
			guard scanner.eat(kind: .punctuation(.colon)) else { return false }
			return true
		} commit: {
			$0.typedBinder()
		}
	}

	static var programParam: Production {
		Production(name: "programParam", kind: .typedBinder) {
			$0.isAt(kind: .punctuation(.lparen))
		} commit: {
			$0.typedBinder()
		}
	}

	static var holeExpression: Production {
		token(
			kind: .punctuation(.questionMark),
			metadata: TokenMetadata(semanticTokenType: .comment, semanticTokenModifiers: [.abstract])
		)
		.wrap(name: "holeExpression", kind: .hole)
	}

	static var universeCodesExpression: Production {
		token(kind: .keyword(.type), metadata: TokenMetadata(semanticTokenType: .type))
			.wrap(name: "universeCodesExpression", kind: .universeCodes)
	}

	static var parenthesisedExpression: Production {
		token(kind: .punctuation(.lparen))
			.then { state in
				_ = expression.precededByTrivia.tryParse(&state, error: "Expected expression")
				_ = token(kind: .punctuation(.rparen)).precededByTrivia.parse(&state, error: "Expected right parenthesis")
			}
	}

	static var reflExpression: Production {
		token(kind: .keyword(.refl), metadata: TokenMetadata(semanticTokenType: .keyword))
			.wrap(name: "reflExpression", kind: .refl)
	}

	static var identifierExpression: Production {
		token(kind: .identifier)
			.wrap(name: "identifierExpression", kind: .identifier)
	}

	static var atomicExpression: Production {
		.cascade([
			recordType,
			holeExpression,
			universeCodesExpression,
			identifierExpression,
			parenthesisedExpression,
			recordExpression,
			reflExpression
		])
	}

	static var depFunType: Production {
		let typedBinder = Self.funTypeBinder
		return typedBinder.then { state in
			while typedBinder.precededByTrivia.tryParse(&state) {}

			let gotArrow =
				token(kind: .punctuation(.rightArrow), metadata: TokenMetadata(semanticTokenType: .keyword))
				.precededByTrivia
				.tryParse(&state, error: "Expected function arrow `->`")

			if gotArrow {
				_ = expression.precededByTrivia.parse(&state, error: "Expected type expression to the right of function arrow")
			}

		}
		.wrap(name: "depFunType", kind: .depFunType)
	}

	static var recordType: Production {
		token(kind: .keyword(.record), metadata: TokenMetadata(semanticTokenType: .keyword)).then { state in
			_ = recordItemSpec.delimitedList(.curly).precededByTrivia.parse(&state, error: "Expected record type specification")
		}
		.wrap(name: "recordType", kind: .recordType)
	}

	private static func afterIdentifier(_ kind: TokenKind) -> (Scanner) -> Bool {
		return { scanner in
			var scanner = scanner
			guard scanner.eat(kind: .identifier) else { return false }
			scanner.eatTrivia()
			return scanner.eat(kind: kind)
		}
	}

	static var recordItemSpec: Production {
		Production(
			name: "recordItemSpec",
			kind: .recordItem(kind: .spec),
			precondition: afterIdentifier(.punctuation(.colon))
		) { state in
			state.expect(kind: .identifier)
			_ = token(kind: .punctuation(.colon), metadata: TokenMetadata(semanticTokenType: .keyword))
				.precededByTrivia
				.parse(&state, error: "Expected `:`")
			_ = expression.precededByTrivia.parse(&state, error: "Expected expression")
		}
	}

	static var namedRecordItemImpl: Production {
		Production(
			name: "recordItemImpl",
			kind: .recordItem(kind: .impl),
			precondition: afterIdentifier(.punctuation(.doubleRightArrow))
		) { state in
			state.expect(kind: .identifier)
			_ = token(kind: .punctuation(.doubleRightArrow), metadata: TokenMetadata(semanticTokenType: .keyword))
				.precededByTrivia
				.parse(&state, error: "Expected `=>`")
			_ = expression.precededByTrivia.parse(&state, error: "Expected expression")
		}
	}

	static var unnamedRecordItemImpl: Production {
		expression.wrap(name: "unnamedRecordItemImpl", kind: .recordItem(kind: .impl))
	}

	static var recordItemImpl: Production {
		.cascade([namedRecordItemImpl, unnamedRecordItemImpl])
	}

	static var recordExpression: Production {
		recordItemImpl.delimitedList(.curly).wrap(name: "recordExpression", kind: .record)
	}

	static var appFrame: Production {
		atomicExpression.wrap(name: "appFrame", kind: .appFrame)
	}

	static var dotFrame: Production {
		token(kind: .punctuation(.dot))
			.then { $0.expect(kind: .identifier) }
			.wrap(name: "dotFrame", kind: .projFrame)
	}

	static var spineFrame: Production {
		.cascade([
			appFrame,
			dotFrame
		])
	}

	private static var operatorTable: [(TokenKind, SyntaxTreeKind, PrecedenceGroup)] {
		[
			(.punctuation(.equal), .idType, .equality),
			(.punctuation(.rightArrow), .funType, .arrow)
		]
	}

	private static func peekOperator(_ scanner: Scanner) -> (TokenKind, SyntaxTreeKind, PrecedenceGroup)? {
		var scanner = scanner
		scanner.eatTrivia()
		for (tokenKind, treeKind, group) in operatorTable {
			if scanner.isAt(kind: tokenKind) {
				return (tokenKind, treeKind, group)
			}
		}
		return nil
	}

	static func spineExpression() -> Production {
		Production(name: "spineExpression") {
			atomicExpression.precondition($0)
		} commit: { state in
			let checkPoint = state.checkpoint
			_ = atomicExpression.parse(&state, error: "Expected expression")

			if spineFrame.precededByTrivia.precondition(state.scanner) {
				state.openAtCheckpoint(checkPoint)
				while spineFrame.precededByTrivia.tryParse(&state) {}
				state.close(kind: .spine)
			}
		}
	}

	static func operatorExpression(lowerBound: PrecedenceGroup?) -> Production {
		Production(name: "operatorExpression") {
			spineExpression().precondition($0)
		} commit: { state in
			var checkPoint = state.checkpoint
			_ = spineExpression().parse(&state, error: "Expected expression")

			while let (tokenKind, treeKind, group) = peekOperator(state.scanner), group.canBind(above: lowerBound) {
				state.openAtCheckpoint(checkPoint)
				state.eatTrivia()
				state.expect(kind: tokenKind, metadata: TokenMetadata(semanticTokenType: .operator))
				_ = Self.operatorExpression(lowerBound: group)
					.precededByTrivia
					.parse(&state, error: "Expected argument on right hand side of `\(tokenKind)` operator")
				state.close(kind: treeKind)
				checkPoint = state.checkpoint
			}
		}
	}

	static var operatorExpression: Production {
		operatorExpression(lowerBound: nil)
	}

	static var expression: Production {
		.cascade([depFunType, lambdaExpression, operatorExpression])
	}

	static var lambdaExpression: Production {
		token(kind: .punctuation(.backslash)).then { state in
			_ = token(kind: .identifier, metadata: TokenMetadata(semanticTokenType: .variable)).wrap(name: "untypedBinder", kind: .untypedBinder).precededByTrivia.many1.tryParse(&state)
			_ = token(kind: .punctuation(.doubleRightArrow)).precededByTrivia.parse(&state, error: "Expected `=>`")
			_ = expression.precededByTrivia.parse(&state, error: "Expected expression")
		}
		.wrap(name: "lambdaExpression", kind: .lambda)
	}

	static var decl: Production {
		.cascade([theory, programDecl])
	}

	static var theory: Production {
		token(kind: .keyword(.theory), metadata: TokenMetadata(semanticTokenType: .keyword))
			.then { state in
				_ = theoryName.precededByTrivia.parse(&state, error: "Expected theory name")
				let gotWhere =
					token(kind: .keyword(.where), metadata: TokenMetadata(semanticTokenType: .keyword))
					.precededByTrivia
					.parse(&state, error: "Expected `where` keyword")
				if gotWhere {
					_ = decl.delimitedList(.block).parse(&state, error: "Expected theory block")
				}
			}
			.wrap(name: "theory", kind: .theory, metadata: SyntaxTreeMetadata(delimitedFoldingRangeKind: .region))
	}

	static var theoryName: Production {
		token(kind: .identifier, metadata: TokenMetadata(semanticTokenType: .interface))
			.wrap(name: "theoryName", kind: .theoryName)
	}

	static var programName: Production {
		token(kind: .identifier)
			.wrap(name: "programName", kind: .programName)
	}

	static var programClauseLhs: Production {
		expression.wrap(name: "programClauseLhs", kind: .programClauseLhs)
	}

	static var programReturn: Production {
		token(kind: .punctuation(.doubleRightArrow), metadata: TokenMetadata(semanticTokenType: .operator))
			.then { state in
				_ = expression.precededByTrivia.tryParse(&state, error: "Expected right-hand side to return expression")
			}
			.wrap(name: "programReturn", kind: .programReturn)
	}

	static var programGadgetSplit: Production {
		token(kind: .keyword(.split))
			.wrap(name: "programGadgetSplit", kind: .splitGadget)
	}

	static var programGadgetIntro: Production {
		token(kind: .keyword(.intro))
			.wrap(name: "programGadgetIntro", kind: .introGadget)
	}

	static var programGadget: Production {
		.cascade([programGadgetSplit, programGadgetIntro])
	}

	static var programBy: Production {
		token(kind: .punctuation(.doubleLeftArrow), metadata: TokenMetadata(semanticTokenType: .operator))
			.then { state in
				_ = programGadget.precededByTrivia.tryParse(&state, error: "Expected `intro` or `split` gadget after `<=`.")
			}
			.wrap(name: "programBy", kind: .programBy)
	}

	static var programCommand: Production {
		.cascade([programReturn, programBy])
	}

	static var programType: Production {
		token(kind: .punctuation(.colon))
			.then { state in
				_ = expression.precededByTrivia.parse(&state, error: "Expected a type expression here")
			}
			.wrap(name: "programType", kind: .programType)
	}

	static var programMatchClause: Production {
		programClauseLhs
			.then { state in
				_ =
					programCommand
					.precededByTrivia
					.wrap(name: "programMatchClauseRhs", kind: .programClauseRhs)
					.tryParse(&state, error: "Expected right-hand side to program match clause")
			}
			.wrap(name: "programMatchClauseRhs", kind: .programMatchClause)
	}

	static var programDeclWhereBody: Production {
		token(kind: .keyword(.where), metadata: TokenMetadata(semanticTokenType: .keyword))
			.then { state in
				_ =
					Production.cascade([programCommand, programMatchClause])
					.delimitedList(.block)
					.precededByTrivia
					.tryParse(&state)
			}
	}

	static var programDeclBody: Production {
		.cascade([programReturn, programDeclWhereBody])
			.wrap(name: "programDeclBody", kind: .programBody)
	}

	static var programDecl: Production {
		programName.then { state in
			let gotParams =
				programParam
				.many1
				.wrap(name: "programParams", kind: .programParams)
				.precededByTrivia
				.tryParse(&state)

			let gotType = programType.precededByTrivia.parse(
				&state,
				error: gotParams ? "Program declaration: expected type annotation here" : "Program declaration: expected parameters or type annotation here"
			)

			_ = programDeclBody.precededByTrivia.tryParse(&state, error: gotType ? "Expected `=>` or `where` clause" : nil)
		}
		.wrap(name: "programDecl", kind: .programDecl, metadata: SyntaxTreeMetadata(delimitedFoldingRangeKind: .region))
	}

	static var importName: Production {
		token(kind: .identifier)
			.wrap(name: "importName", kind: .importName)
	}

	static var importDecl: Production {
		token(kind: .keyword(.import), metadata: TokenMetadata(semanticTokenType: .keyword))
			.then { state in
				_ = importName.precededByTrivia.parse(&state, error: "Expected import name")  // TODO
			}
			.wrap(name: "importDecl", kind: .import)
	}

	static var documentItem: Production {
		.cascade([importDecl, decl])
	}

	static var document: Production {
		let documentBlockAndEof =
			documentItem.delimitedList(.block).then { state in
				_ = token(kind: .eof).precededByTrivia.parse(&state, error: "Expected EOF")
			}

		return
			.cascade([token(kind: .eof), documentBlockAndEof])
			.precededByTrivia
			.wrap(name: "document", kind: .document)
	}
}

extension Parser {
	public mutating func document() {
		_ = Self.document.parse(&self, error: "Invalid start of document")
	}
}