roc/benchmarks/AStar.roc

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