Roguelike Dungeon Generation

main.nim

import math, random, strformat, tables

# Original Python version: https://github.com/DanaL/RLDungeonGenerator

type
  DungeonSqr = object
    sqr: string
  Room = object
    row: int
    col: int
    height: int
    width: int
  RLDungeonGenerator = object
    MAX: int
    width: int
    height: int
    leaves: seq[tuple[minRow: int, minCol: int, maxRow: int, maxCol: int]]
    # leaves: seq[tuple[int, int, int, int]]
    # leaves: seq[(int, int, int, int)]
    dungeon: seq[seq[DungeonSqr]]
    rooms: seq[Room]


proc splitOnVertical(self: var RLDungeonGenerator; minRow, minCol, maxRow, maxCol: int)
proc splitOnHorizontal(self: var RLDungeonGenerator; minRow, minCol, maxRow, maxCol: int)

proc getCh(self: DungeonSqr): string =
  self.sqr

proc initRLDungeonGenerator(width: int, height: int): RLDungeonGenerator  =
  result.MAX = 15 # Cutoff for when we want to stop dividing sections
  result.width = width
  result.height = height
  var i = 0
  while i < result.height:
    var j = 0
    var row: seq[DungeonSqr]
    while j < result.width:
      row.add DungeonSqr(sqr: "#")
      j += 1
    i += 1
    result.dungeon.add move(row)

proc printMap(self: RLDungeonGenerator) =
  for r in 0..<self.height:
    var row = ""
    for c in 0..<self.width:
      row.add self.dungeon[r][c].getCh()
    echo row

proc randomSplit(self: var RLDungeonGenerator; minRow, minCol, maxRow, maxCol: int) =
  # We want to keep splitting until the sections get down to the threshold
  let seg_height = maxRow - minRow
  let seg_width = maxCol - minCol
  # echo $seg_height & " " & $seg_width
  if seg_height < self.MAX and seg_width < self.MAX:
    self.leaves.add (minRow, minCol, maxRow, maxCol)
  elif seg_height < self.MAX and seg_width >= self.MAX:
    self.splitOnVertical(minRow, minCol, maxRow, maxCol)
  elif seg_height >= self.MAX and seg_width < self.MAX:
    self.splitOnHorizontal(minRow, minCol, maxRow, maxCol)
  else:
    if rand(1.0) < 0.5:
      self.splitOnHorizontal(minRow, minCol, maxRow, maxCol)
    else:
      self.splitOnVertical(minRow, minCol, maxRow, maxCol)

proc splitOnHorizontal(self: var RLDungeonGenerator, minRow, minCol, maxRow, maxCol: int) =
    let split = (minRow + maxRow) div 2 + sample([-2, -1, 0, 1, 2])
    self.randomSplit(minRow, minCol, split, maxCol)
    self.randomSplit(split + 1, minCol, maxRow, maxCol)

proc splitOnVertical(self: var RLDungeonGenerator, minRow, minCol, maxRow, maxCol: int) =
    let split = (minCol + maxCol) div 2 + sample([-2, -1, 0, 1, 2])
    self.randomSplit(minRow, minCol, maxRow, split)
    self.randomSplit(minRow, split + 1, maxRow, maxCol)


proc carve_rooms(self: var RLDungeonGenerator) =
  for leaf in self.leaves:
    # We don't want to fill in every possible room or the 
    # dungeon looks too uniform
    if rand(1.0) > 0.80:
      continue
    let sectionWidth: int = leaf.maxCol - leaf.minCol
    let sectionHeight: int = leaf.maxRow - leaf.minRow

    # The actual room's height and width will be 60-100% of the 
    # available section. 
    let roomWidth = round(rand(60 ..< 100) / 100 * sectionWidth.float).int
    let roomHeight = round(rand(60 ..< 100) / 100 * sectionHeight.float).int

    # If the room doesn't occupy the entire section we are carving it from,
    # 'jiggle' it a bit in the square
    let roomStartRow =
      if sectionHeight > roomHeight:
        leaf[0] + rand(0 ..< sectionHeight - roomHeight)
      else:
        leaf[0]
    let roomStartCol =
      if sectionWidth > roomWidth:
        leaf[1] + rand(0 ..< sectionWidth - roomWidth)
      else:
        leaf[1]

    self.rooms.add(
        Room(row: roomStartRow, col: roomStartCol, height: roomHeight, width: roomWidth))
    for r in roomStartRow..<roomStartRow + roomHeight:
      for c in roomStartCol..<roomStartCol + roomWidth:
        self.dungeon[r][c] = DungeonSqr(sqr: ".")

proc are_rooms_adjacent(self: RLDungeonGenerator; room1, room2: Room): (seq[int], seq[int]) =
  var adj_rows: seq[int]
  var adj_cols: seq[int]
  for r in room1.row ..< room1.row + room1.height:
    if r >= room2.row and r < room2.row + room2.height:
      adj_rows.add r
  for c in room1.col ..< room1.col + room1.width:
    if c >= room2.col and c < room2.col + room2.width:
      adj_cols.add c
  result = (adj_rows, adj_cols)

proc distance_between_rooms(self: RLDungeonGenerator; room1, room2: Room ): float =
  let centre1 = (room1.row + room1.height div 2, room1.col + room1.width div 2)
  let centre2 = (room2.row + room2.height div 2, room2.col + room2.width div 2)
  result = sqrt(float((centre1[0] - centre2[0]) ^ 2 + (centre1[1] - centre2[1]) ^ 2))

proc carve_corridor_between_rooms(self: var RLDungeonGenerator; room1: Room, room2: (Room, seq[int], string, float)) =
  # echo "in"
  if room2[2] == "rows":
    let row = sample(room2[1])
    # Figure out which room is to the left of the other
    let (start_col, end_col) =
      if room1.col + room1.width < room2[0].col:
        (room1.col + room1.width, room2[0].col)
      else:
        (room2[0].col + room2[0].width, room1.col)

    for c in start_col ..< end_col:
      self.dungeon[row][c] = DungeonSqr(sqr: ".")

    if end_col - start_col >= 4:
      self.dungeon[row][start_col] = DungeonSqr(sqr: "+")
      self.dungeon[row][end_col - 1] = DungeonSqr(sqr: "+")
    elif start_col == end_col - 1:
      self.dungeon[row][start_col] = DungeonSqr(sqr: "+")
  else:
    let col = sample(room2[1])
    # Figure out which room is above the other
    let (start_row, end_row) =
      if room1.row + room1.height < room2[0].row:
        (room1.row + room1.height, room2[0].row)
      else:
        (room2[0].row + room2[0].height, room1.row)

    for r in start_row ..< end_row:
      self.dungeon[r][col] = DungeonSqr(sqr: ".")

    if end_row - start_row >= 4:
      self.dungeon[start_row][col] = DungeonSqr(sqr: "+")
      self.dungeon[end_row - 1][col] = DungeonSqr(sqr: "+")
    elif start_row == end_row - 1:
      self.dungeon[start_row][col] = DungeonSqr(sqr: "+")

proc find_closest_unconnect_groups(self: var RLDungeonGenerator,
    groups: var seq[seq[Room]],
    room_dict: Table[(int, int), seq[(Room, seq[int], string, float)]]) =

  var shortest_distance = 99999.0
  var start: Room
  # var start_group: ptr seq[Room]
  var start_group: int
  var nearest: (Room, seq[int], string, float)

  for i, group in groups:
    for room in group:
      let key = (room.row, room.col)
      for other in room_dict[key]:
        if other[0] notin group and other[3] < shortest_distance:
          shortest_distance = other[3]
          start = room
          nearest = other
          # start_group = groups[i].addr
          start_group = i

  self.carve_corridor_between_rooms(start, nearest)

  # Merge the groups
  var other_group: seq[Room]
  for group in groups:
    if nearest[0] in group:
      other_group = group
      break

  # start_group[].add(other_group)
  groups[start_group].add(other_group)
  let removeIndex = groups.find(other_group)
  if removeIndex != -1:
    groups.delete(removeIndex)

proc connectRooms(self: var RLDungeonGenerator) =
  # Build a dictionary containing an entry for each room. Each bucket will
  # hold a list of the adjacent rooms, weather they are adjacent along rows or 
  # columns and the distance between them.
  #
  # Also build the initial groups (which start of as a list of individual rooms)
  var groups: seq[seq[Room]]
  var room_dict = initTable[(int, int), seq[(Room, seq[int], string, float)]]()

  for room in self.rooms:
    let key = (room.row, room.col)
    room_dict[key] = @[]
    for other in self.rooms:
      let other_key = (other.row, other.col)
      if key == other_key:
        continue
      let adj = self.are_rooms_adjacent(room, other)
      if adj[0].len > 0:
        room_dict[key].add((other, adj[0], "rows", self.distance_between_rooms(room, other)))
      elif adj[1].len > 0:
        room_dict[key].add((other, adj[1], "cols", self.distance_between_rooms(room, other)))

    groups.add @[room]

  while groups.len > 1:
    self.find_closest_unconnect_groups(groups, room_dict)

proc generateMap(self: var RLDungeonGenerator) =
    self.randomSplit(1, 1, self.height - 1, self.width - 1)
    self.carveRooms()
    self.connectRooms()

proc main() =
  randomize()
  # var dg = initRLDungeonGenerator(width = 5, height = 5)
  # var dg = initRLDungeonGenerator(width = 30, height = 30)
  var dg = initRLDungeonGenerator(width = 75, height = 40)
  dg.generateMap()
  dg.printMap()

when isMainModule:
  main()

main.py

#!/usr/bin/env python3

from math import sqrt
from random import random
from random import randrange
from random import choice
from random import shuffle
import time
import os

clear = r"printf '\33c\e[3J\33c'"

Red = "\033[31m"
Green = "\033[32m"
Yellow = "\033[33m"
Blue = "\033[34m"
Magenta = "\033[35m"
Cyan = "\033[36m"
LightGray = "\033[37m"
DarkGray = "\033[90m"
LightRed = "\033[91m"
LightGreen = "\033[92m"
LightYellow = "\033[93m"
LightBlue = "\033[94m"
LightMagenta = "\033[95m"
LightCyan = "\033[96m"

color_end = "\033[0m"
color_index = 0
color = [
    Red,
    Green,
    Yellow,
    Blue,
    Magenta,
    Cyan,
    LightRed,
    LightGreen,
    LightYellow,
    LightBlue,
    LightMagenta,
    LightCyan,
]


class DungeonSqr:
  def __init__(self, sqr):
    self.sqr = sqr

  def get_ch(self):
    return self.sqr

class Room:
  def __init__(self, r, c, h, w):
    self.row = r
    self.col = c
    self.height = h
    self.width = w

class RLDungeonGenerator:
  def __init__(self, w, h):
    # self.MAX = 20  # Cutoff for when we want to stop dividing sections
    self.MAX = 15  # Cutoff for when we want to stop dividing sections
    # self.MAX = 20  # Cutoff for when we want to stop dividing sections
    self.width = w
    self.height = h
    self.leaves = []
    self.dungeon = []
    self.rooms = []
    self.state = []
    # self.state = [["#" for j in range(w)] for i in range(h)]
    for i in range(self.height):
      row = []
      for j in range(self.width):
        row.append("#")
      self.state.append(row)
    # print(self.state)
    # print(len(self.state))
    # quit(-1)

    for h in range(self.height):
      row = []
      for w in range(self.width):
        row.append(DungeonSqr("#"))
      self.dungeon.append(row)

  def echo(self, min_row, min_col, max_row, max_col):
    global color
    global color_index

    if len(color) == color_index:
      shuffle(color)
      color_index = 0

    for i in range(self.height):
      for j in range(self.width):
        if (min_row <= i and i < max_row) and (min_col <= j and j < max_col):
          self.state[i][j] = f"{color[color_index]}#{color_end}"
      print("".join(self.state[i]))
    time.sleep(0.15)
    os.system(clear)
    color_index = color_index + 1

  def random_split(self, min_row, min_col, max_row, max_col):
    # We want to keep splitting until the sections get down to the threshold
    seg_height = max_row - min_row
    seg_width = max_col - min_col

    if seg_height < self.MAX and seg_width < self.MAX:
      self.leaves.append((min_row, min_col, max_row, max_col))
      # print(max_col - min_col)
      self.echo(min_row, min_col, max_row, max_col)
    elif seg_height < self.MAX and seg_width >= self.MAX:
      # 垂直
      self.split_on_vertical(min_row, min_col, max_row, max_col)
    elif seg_height >= self.MAX and seg_width < self.MAX:
      # 水平
      self.split_on_horizontal(min_row, min_col, max_row, max_col)
    else:
      if random() < 0.5:
        self.split_on_horizontal(min_row, min_col, max_row, max_col)
      else:
        self.split_on_vertical(min_row, min_col, max_row, max_col)

  def split_on_horizontal(self, min_row, min_col, max_row, max_col):
    split = (min_row + max_row) // 2 + choice((-2, -1, 0, 1, 2))
    self.random_split(min_row, min_col, split, max_col)
    self.random_split(split + 1, min_col, max_row, max_col)

  def split_on_vertical(self, min_row, min_col, max_row, max_col):
    split = (min_col + max_col) // 2 + choice((-2, -1, 0, 1, 2))
    self.random_split(min_row, min_col, max_row, split)
    self.random_split(min_row, split + 1, max_row, max_col)

  def carve_rooms(self):
    for leaf in self.leaves:
      # We don't want to fill in every possible room or the
      # dungeon looks too uniform
      if random() > 0.80:
        continue
      section_width = leaf[3] - leaf[1]  # max_col - min_col
      section_height = leaf[2] - leaf[0]  # max_row - min_row

      # The actual room's height and width will be 60-100% of the
      # available section.
      room_width = round(randrange(60, 100) / 100 * section_width)
      room_height = round(randrange(60, 100) / 100 * section_height)

      # If the room doesn't occupy the entire section we are carving it from,
      # 'jiggle' it a bit in the square

      if section_height > room_height:
        room_start_row = leaf[0] + randrange(section_height - room_height)
      else:
        room_start_row = leaf[0]

      if section_width > room_width:
        room_start_col = leaf[1] + randrange(section_width - room_width)
      else:
        room_start_col = leaf[1]

      # room_start_col = leaf[1]
      # room_start_row = leaf[0]

      self.rooms.append(
        Room(room_start_row, room_start_col, room_height, room_width)
      )
      for r in range(room_start_row, room_start_row + room_height):
        for c in range(room_start_col, room_start_col + room_width):
          self.dungeon[r][c] = DungeonSqr(".")

  def are_rooms_adjacent(self, room1, room2):
    adj_rows = []
    adj_cols = []
    for r in range(room1.row, room1.row + room1.height):
      if r >= room2.row and r < room2.row + room2.height:
        adj_rows.append(r)

    for c in range(room1.col, room1.col + room1.width):
      if c >= room2.col and c < room2.col + room2.width:
        adj_cols.append(c)

    return (adj_rows, adj_cols)

  def distance_between_rooms(self, room1, room2):
    centre1 = (room1.row + room1.height // 2, room1.col + room1.width // 2)
    centre2 = (room2.row + room2.height // 2, room2.col + room2.width // 2)

    return sqrt((centre1[0] - centre2[0]) ** 2 + (centre1[1] - centre2[1]) ** 2)

  def carve_corridor_between_rooms(self, room1, room2):
    # print(room1 == None)
    # print(room2 == None)
    if room2[2] == "rows":
      row = choice(room2[1])
      # Figure out which room is to the left of the other
      if room1.col + room1.width < room2[0].col:
        start_col = room1.col + room1.width
        end_col = room2[0].col
      else:
        start_col = room2[0].col + room2[0].width
        end_col = room1.col
      for c in range(start_col, end_col):
        self.dungeon[row][c] = DungeonSqr(".")

      if end_col - start_col >= 4:
        self.dungeon[row][start_col] = DungeonSqr("+")
        self.dungeon[row][end_col - 1] = DungeonSqr("+")
      elif start_col == end_col - 1:
        self.dungeon[row][start_col] = DungeonSqr("+")
    else:
      col = choice(room2[1])
      # Figure out which room is above the other
      if room1.row + room1.height < room2[0].row:
        start_row = room1.row + room1.height
        end_row = room2[0].row
      else:
        start_row = room2[0].row + room2[0].height
        end_row = room1.row

      for r in range(start_row, end_row):
        self.dungeon[r][col] = DungeonSqr(".")

      if end_row - start_row >= 4:
        self.dungeon[start_row][col] = DungeonSqr("+")
        self.dungeon[end_row - 1][col] = DungeonSqr("+")
      elif start_row == end_row - 1:
        self.dungeon[start_row][col] = DungeonSqr("+")

    # Find two nearby rooms that are in difference groups, draw
    # a corridor between them and merge the groups
  def find_closest_unconnect_groups(self, groups, room_dict):
    shortest_distance = 99999
    start = None
    start_group = None
    nearest = None

    for group in groups:
      for room in group:
        key = (room.row, room.col)
        for other in room_dict[key]:
          if not other[0] in group and other[3] < shortest_distance:
            shortest_distance = other[3]
            start = room
            nearest = other
            start_group = group

    # if nearest == None:
    #   print(groups)
    #   self.print_map()
    #   quit(9)

    # self.print_map()
    self.carve_corridor_between_rooms(start, nearest)

    # Merge the groups
    other_group = None
    for group in groups:
      if nearest[0] in group:
        other_group = group
        break

    start_group += other_group
    groups.remove(other_group)

  def connect_rooms(self):
    # Build a dictionary containing an entry for each room. Each bucket will
    # hold a list of the adjacent rooms, weather they are adjacent along rows or
    # columns and the distance between them.
    #
    # Also build the initial groups (which start of as a list of individual rooms)
    groups = []
    room_dict = {}
    for room in self.rooms:
      key = (room.row, room.col)
      room_dict[key] = []
      for other in self.rooms:
        other_key = (other.row, other.col)
        if key == other_key:
          continue
        adj = self.are_rooms_adjacent(room, other)
        if len(adj[0]) > 0:
          room_dict[key].append(
              (
                other,
                adj[0],
                "rows",
                self.distance_between_rooms(room, other),
                )
              )
        elif len(adj[1]) > 0:
          room_dict[key].append(
              (
                other,
                adj[1],
                "cols",
                self.distance_between_rooms(room, other),
                )
              )

      groups.append([room])

    while len(groups) > 1:
      self.find_closest_unconnect_groups(groups, room_dict)

  def generate_map(self):
    self.random_split(1, 1, self.height - 1, self.width - 1)
    self.carve_rooms()
    self.connect_rooms()

  def print_map2(self):
    for i in range(self.height):
      for j in range(self.width):
        if self.dungeon[i][j].get_ch() == ".":
          self.state[i][j] = "."
        elif self.dungeon[i][j].get_ch() == "+":
          self.state[i][j] = "+"
      print("".join(self.state[i]))

  def print_map(self):
    for r in range(self.height):
      row = ""
      for c in range(self.width):
        row += self.dungeon[r][c].get_ch()
      print(row)

shuffle(color)
os.system(clear)
# print(f"{bcolors.WARNING}Warning: No active frommets remain. Continue?{bcolors.ENDC}")
# dg = RLDungeonGenerator(75, 40)
dg = RLDungeonGenerator(60, 30)
# dg = RLDungeonGenerator(20, 20)
dg.generate_map()
dg.print_map2()
# dg.print_map()