roc/examples/benchmarks/AStar.roc
2021-03-14 21:57:10 -04:00

137 lines
3.7 KiB
Plaintext

interface AStar exposes [ findPath, Model, initialModel, cheapestOpen, takeStep, reconstructPath ] imports [Quicksort]
findPath = \costFn, moveFn, start, end ->
astar costFn moveFn end (initialModel start)
Model position :
{
evaluated : Set position,
openSet : Set position,
costs : Dict position F64,
cameFrom : Dict position position
}
initialModel : position -> Model position
initialModel = \start ->
{
evaluated : Set.empty,
openSet : Set.single start,
costs : Dict.single start 0,
cameFrom : Dict.empty
}
cheapestOpen : (position -> F64), Model position -> Result position {}
cheapestOpen = \costFn, model ->
model.openSet
|> Set.toList
|> List.keepOks (\position ->
when Dict.get model.costs position is
Err _ ->
Err {}
Ok cost ->
Ok { cost: cost + costFn position, position }
)
|> Quicksort.sortBy .cost
|> List.first
|> Result.map .position
|> Result.mapErr (\_ -> {})
reconstructPath : Dict position position, position -> List position
reconstructPath = \cameFrom, goal ->
when Dict.get cameFrom goal is
Err _ ->
[]
Ok next ->
List.append (reconstructPath cameFrom next) goal
updateCost : position, position, Model position -> Model position
updateCost = \current, neighbor, model ->
newCameFrom =
Dict.insert model.cameFrom neighbor current
newCosts =
Dict.insert model.costs neighbor distanceTo
distanceTo =
reconstructPath newCameFrom neighbor
|> List.len
|> Num.toFloat
newModel =
{ model &
costs: newCosts,
cameFrom: newCameFrom
}
when Dict.get model.costs neighbor is
Err _ ->
newModel
Ok previousDistance ->
if distanceTo < previousDistance then
newModel
else
model
astar : (position, position -> F64), (position -> Set position), position, Model position -> Result (List position) {}
astar = \costFn, moveFn, goal, model ->
when cheapestOpen (\source -> costFn source goal) model is
Err {} ->
Err {}
Ok current ->
if current == goal then
Ok (reconstructPath model.cameFrom goal)
else
modelPopped =
{ model &
openSet: Set.remove model.openSet current,
evaluated: Set.insert model.evaluated current,
}
neighbors =
moveFn current
newNeighbors =
Set.difference neighbors modelPopped.evaluated
modelWithNeighbors =
{ modelPopped &
openSet: Set.union modelPopped.openSet newNeighbors
}
modelWithCosts =
Set.walk newNeighbors (\n, m -> updateCost current n m) modelWithNeighbors
astar costFn moveFn goal modelWithCosts
takeStep = \moveFn, _goal, model, current ->
modelPopped =
{ model &
openSet: Set.remove model.openSet current,
evaluated: Set.insert model.evaluated current,
}
neighbors = moveFn current
newNeighbors = Set.difference neighbors modelPopped.evaluated
modelWithNeighbors = { modelPopped & openSet: Set.union modelPopped.openSet newNeighbors }
# a lot goes wrong here
modelWithCosts =
Set.walk newNeighbors (\n, m -> updateCost current n m) modelWithNeighbors
modelWithCosts