-- N-Queens problem

def fourQueens : [[Integer]] := matchAll [1,2,3,4] as multiset integer with
  | $a_1 ::
     (!#(a_1 - 1) & !#(a_1 + 1) & $a_2) ::
      (!#(a_1 - 2) & !#(a_1 + 2) & !#(a_2 - 1) & !#(a_2 + 1) & $a_3) ::
       (!#(a_1 - 3) & !#(a_1 + 3) & !#(a_2 - 2) & !#(a_2 + 2) & !#(a_3 - 1) & !#(a_3 + 1) & $a_4) ::
        []
   -> [a_1,a_2,a_3,a_4]

assertEqual "four queens"
  fourQueens
  [[2,4,1,3],[3,1,4,2]]

def nQueens (n: Integer) : [[Integer]] := matchAll [1..n] as multiset integer with
  | $a_1 ::
      (loop $i (2, n)
         ((loop $j (1, i - 1)
             (!#(a_j - (i - j)) & !#(a_j + (i - j)) & ...)
             $a_i) :: ...)
         [])
  -> map (\i -> a_i) [1..n]

assertEqual "nQueens 4"
  (nQueens 4)
  [[2,4,1,3],[3,1,4,2]]

def fourQueens2 := matchAll [1,2,3,4] as multiset integer with
  | $a_1 ::
     (!#(a_1 - 1) & !#(a_1 + 1) & $a_2) ::
      (!#(a_1 - 2) & !#(a_1 + 2) & !#(a_2 - 1) & !#(a_2 + 1) & $a_3) ::
       (!#(a_1 - 3) & !#(a_1 + 3) & !#(a_2 - 2) & !#(a_2 + 2) & !#(a_3 - 1) & !#(a_3 + 1) & $a_4) ::
        []
   -> a

assertEqual "four queens 2 (hash result)"
  (length fourQueens2)
  2