toysolver-0.5.0: samples/programs/numberlink/numberlink.hs
{-# LANGUAGE FlexibleContexts, OverloadedStrings #-}
{-# OPTIONS_GHC -Wall #-}
module Main where
import Control.Applicative hiding (many, optional)
import Control.Monad
import Data.Array.IArray
import qualified Data.ByteString.Lazy.Char8 as BL
import Data.Char
import Data.Default.Class
import Data.List
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe
import qualified Data.PseudoBoolean as PB
import Data.Set (Set)
import qualified Data.Set as Set
import System.Console.GetOpt
import System.Environment
import System.Exit
import System.IO
import Text.Parsec hiding (try)
import qualified Text.Parsec.ByteString.Lazy as ParsecBL
import qualified ToySolver.SAT as SAT
import qualified ToySolver.SAT.PBO as PBO
import qualified ToySolver.SAT.Store.PB as PBStore
data Problem
= Problem
{ probSize :: (Int,Int,Int)
, probIsMultiLayer :: Bool
, probLineNum :: Int
, probLines :: [(Number, Cell, Cell)]
, probVias :: [(Via, [Cell])]
}
deriving (Show)
type Number = Int
type Cell = (Int,Int,Int)
type Edge = (Cell, Cell)
type Via = String
parser :: Stream s m Char => ParsecT s u m Problem
parser = do
optional $ string "\xEF\xBB\xBF" -- BOM in UTF-8
spaces
_ <- string "SIZE"
spaces
w <- num
_ <- char 'X'
h <- num
dd <- optionMaybe $ char 'X' >> num
let (d,isMultiLayer) =
case dd of
Nothing -> (1, False)
Just d -> (d, True)
_ <- endOfLine
let cell = do
_ <- char '('
x <- num
_ <- char ','
y <- num
z <-
if isMultiLayer then do
_ <- char ','
num
else
return 1
_ <- char ')'
spaces
return (x,y,z)
_ <- string "LINE_NUM"
spaces
n <- num
_ <- endOfLine
xs <- many $ do
_ <- string "LINE#"
i <- num
spaces
src <- cell
optional $ char '-' >> spaces
dst <- cell
return (i,src,dst)
ys <- many $ do
_ <- string "VIA#"
v <- many letter
spaces
cs <- cell `sepBy` optional (char '-' >> spaces)
return (v,cs)
return $ Problem (w,h,d) (isJust dd) n xs ys
where
num = read <$> many digit
type Encoded = (Array Cell (Array Number SAT.Var), Map Edge SAT.Var)
encode :: SAT.AddPBLin m enc => enc -> Options -> Problem -> m Encoded
encode enc opt Problem{ probSize = (w,h,d), probLineNum = n, probLines = ls, probVias = vias } = do
let bnd = ((0,0,1), (w-1,h-1,d))
cells = range bnd
edges = [(a,b) | a@(x,y,z) <- cells, b <- [(x+1,y,z),(x,y+1,z)], inRange bnd b]
adjacentEdges a@(x,y,z) =
[(a,b) | b <- [(x+1,y,z),(x,y+1,z)], inRange bnd b] ++
[(b,a) | b <- [(x-1,y,z),(x,y-1,z)], inRange bnd b]
ks = [1..n]
-- ビアへの数字の割当を表す変数 (0 は数字なしを表す)
viaVs <- liftM Map.fromList $ forM vias $ \(via, _) -> do
let r = if optAssumeNoBlank opt
then (1,n)
else (0,n)
zs <- liftM (array r) $ forM (range r) $ \k -> do
v <- SAT.newVar enc
return (k,v)
return (via, zs)
let viaPos = Map.fromList [(a,via) | (via,as) <- vias, a <- as]
-- 各マスへの数字の割り当てを表す変数 (0 は数字なしを表す)
vs <- liftM (array bnd) $ forM cells $ \a -> do
case Map.lookup a viaPos of
Just via -> return (a, viaVs Map.! via)
Nothing -> do
let r = if optAssumeNoBlank opt
then (1,n)
else (0,n)
zs <- liftM (array r) $ forM (range r) $ \k -> do
v <- SAT.newVar enc
return (k,v)
return (a, zs)
-- 各辺の有無を表す変数
evs <- liftM Map.fromList $ forM edges $ \e -> do
v <- SAT.newVar enc
return (e,v)
-- 初期数字
let m0 = Map.fromList [(c,k) | (k,src,dst) <- ls, c <- [src,dst]]
forM_ (Map.toList m0) $ \(c,k) -> do
SAT.addClause enc [vs!c!k]
-- 各マスには高々ひとつしか数字が入らない
forM_ (range bnd) $ \a -> do
SAT.addExactly enc [v | v <- elems (vs!a)] 1
forM_ (Map.toList evs) $ \((a,b),v) ->
forM_ ks $ \k -> do
-- 辺で連結されるマスは同じ数字
SAT.addClause enc [-v, -(vs!a!k), vs!b!k]
SAT.addClause enc [-v, -(vs!b!k), vs!a!k]
-- 連結されない隣接マスは違う数字
when (optAssumeNoDetour opt) $
SAT.addClause enc [v, -(vs!a!k), -(vs!b!k)]
forM_ (range bnd) $ \a -> do
if a `Map.member` m0 then do
-- 初期数字の入っているマスの次数は1
SAT.addExactly enc [evs Map.! e | e <- adjacentEdges a] 1
else if a `Map.member` viaPos then do
-- ビアの次数は高々1 (ADC2016ルールでも中間層は0になるので注意)
SAT.addAtMost enc [evs Map.! e | e <- adjacentEdges a] 1
else do
-- 初期数字の入っていないマスの次数は0か2
if optAssumeNoBlank opt then do
SAT.addExactly enc [evs Map.! e | e <- adjacentEdges a] 2
else do
let v = vs!a!0
-- v → deg(a)=0
-- SAT.addPBAtMostSoft enc v [(1, evs Map.! e) | e <- adjacentEdges a] 0
forM_ (adjacentEdges a) $ \e -> SAT.addClause enc [-v, -(evs Map.! e)]
-- ¬v → deg(a)=2
SAT.addPBExactlySoft enc (-v) [(1, evs Map.! e) | e <- adjacentEdges a] 2
-- ¬v → deg(a)>0
-- SAT.addClause enc $ v : [evs Map.! e | e <- adjacentEdges a]
-- コの字の禁止
when (optAssumeNoDetour opt) $ do
forM_ [1..d] $ \z -> do
forM_ [0..w-2] $ \x -> do
forM_ [0..h-2] $ \y -> do
let a = (x, y, z)
b = (x+1, y, z)
c = (x, y+1, z)
d = (x+1, y+1, z)
SAT.addAtMost enc [evs Map.! e | e <- [(a,b),(a,c),(b,d),(c,d)]] 2
-- https://kaigi.org/jsai/webprogram/2016/pdf/67.pdf の追加成約
when (optJSAI2016 opt) $ do
let bs = Set.fromList [a | a <- range bnd, a `Map.notMember` m0, a `Map.notMember` viaPos]
forM_ (range bnd) $ \a@(x,y,z) -> do
let a_n = (x,y-1,z)
a_s = (x,y+1,z)
a_w = (x-1,y,z)
a_e = (x+1,y,z)
a_nw = Set.fromList $ takeWhile (inRange bnd) $ tail $ iterate (\(x,y,z) -> (x-1,y-1,z)) a
a_ne = Set.fromList $ takeWhile (inRange bnd) $ tail $ iterate (\(x,y,z) -> (x+1,y-1,z)) a
a_sw = Set.fromList $ takeWhile (inRange bnd) $ tail $ iterate (\(x,y,z) -> (x-1,y+1,z)) a
a_se = Set.fromList $ takeWhile (inRange bnd) $ tail $ iterate (\(x,y,z) -> (x+1,y+1,z)) a
when (inRange bnd a_n && inRange bnd a_w && a_nw `Set.isSubsetOf` bs) $ do
SAT.addAtMost enc [evs Map.! (a_w,a), evs Map.! (a_n,a)] 1
when (inRange bnd a_n && inRange bnd a_e && a_ne `Set.isSubsetOf` bs) $ do
SAT.addAtMost enc [evs Map.! (a,a_e), evs Map.! (a_n,a)] 1
when (inRange bnd a_s && inRange bnd a_w && a_sw `Set.isSubsetOf` bs) $ do
SAT.addAtMost enc [evs Map.! (a_w,a), evs Map.! (a,a_s)] 1
when (inRange bnd a_s && inRange bnd a_e && a_se `Set.isSubsetOf` bs) $ do
SAT.addAtMost enc [evs Map.! (a,a_e), evs Map.! (a,a_s)] 1
forM_ ls $ \(k, (_,_,z1), (_,_,z2)) -> do
let vsk = [vs!k | (_, vs) <- Map.toList viaVs]
if z1 == z2 then do
-- ADC2016では(層をまたぐLINE数)=(VIA数)なので、他のLINEはVIAを使えない
when (optADC2016 opt) $ SAT.addAtMost enc vsk 0
else do
-- 異なる盤面にある数字はビアを必ず使う
SAT.addClause enc vsk
-- ADC2016では(層をまたぐLINE数)=(VIA数)
when (optADC2016 opt) $ SAT.addAtMost enc vsk 1
return (vs, evs)
encodeObj :: SAT.AddPBLin m enc => enc -> Options -> Problem -> Encoded -> m SAT.PBLinSum
encodeObj enc opt Problem{ probSize = (w,h,d) } (cells,edges) = do
let (o1, o2) = fromJust (optOptimize opt)
obj1 <-
if not o1 then do
return []
else if optAssumeNoBlank opt then do
v <- SAT.newVar enc
SAT.addClause enc [v]
return [(fromIntegral (w*h*d), v)]
else
return [(1, -(vs!0)) | vs <- elems cells]
obj2 <-
if not o2 then do
return []
else do
forM (range (bounds cells)) $ \a@(x,y,z) -> do
let w = ((x-1,y,z),a)
e = (a,(x+1,y,z))
n = ((x,y-1,z),a)
s = (a,(x,y+1,z))
v <- SAT.newVar enc
forM_ [e,w] $ \e1 -> do
case Map.lookup e1 edges of
Nothing -> return ()
Just v1 -> do
forM_ [n,s] $ \e2 -> do
case Map.lookup e2 edges of
Nothing -> return ()
Just v2 -> SAT.addClause enc [-v1, -v2, v]
return (1,v)
return $ obj1 ++ obj2
type Solution = (Map Cell Number, Set Edge)
decode :: Problem -> Encoded -> SAT.Model -> Solution
decode prob (vs, evs) m = (solCells, solEdges)
where
solCells = Map.fromList $ do
a <- range (bounds vs)
guard $ a `Set.member` usedCells
case [k | (k,v) <- assocs (vs!a), SAT.evalLit m v] of
k:_ -> return (a,k)
[] -> mzero
solEdges = Set.fromList [e | e@(a,_) <- edges, a `Set.member` usedCells]
edges = [e | (e,ev) <- Map.toList evs, SAT.evalLit m ev]
adjacents = Map.fromListWith Set.union $ concat $ [[(a, Set.singleton b), (b, Set.singleton a)] | (a,b) <- edges]
usedVias = Set.fromList [(x,y,z) | ((x,y),zs) <- Map.toList cs, z <- [Set.findMin zs .. Set.findMax zs]]
where
cs = Map.fromList [((x,y), Set.singleton z) | (_,as) <- probVias prob, a@(x,y,z) <- as, a `Map.member` adjacents]
usedCells = Set.union usedVias cs
where
cs = Set.unions [g (Set.fromList [a,b]) Set.empty | (_k,a,b) <- probLines prob]
g xs visited
| Set.null xs = visited
| otherwise = g (Set.unions [Map.findWithDefault Set.empty x adjacents Set.\\ visited | x <- Set.toList xs]) (Set.union xs visited)
evalObj :: Options -> Problem -> Solution -> Integer
evalObj opt Problem{ probSize = (w,h,d) } (cells,edges) = obj1 + obj2
where
(o1, o2) = fromJust (optOptimize opt)
bnd = ((0,0,1),(w-1,h-1,d))
obj1
| not o1 = 0
| otherwise = fromIntegral $ Map.size cells
obj2
| not o2 = 0
| otherwise = sum $ do
a@(x,y,z) <- range bnd
let w = ((x-1,y,z),a)
e = (a,(x+1,y,z))
n = ((x,y-1,z),a)
s = (a,(x,y+1,z))
if null [() | e1 <- [e,w], e1 `Set.member` edges, e2 <- [n,s], e2 `Set.member` edges] then
return 0
else
return 1
createSolver :: Options -> Problem -> IO (IO (Maybe Solution))
createSolver opt prob = do
solver <- SAT.newSolver
SAT.setLogger solver $ hPutStrLn stderr
encoded <- encode solver opt prob
unless (optAssumeNoBlank opt) $ do
forM_ (elems (fst encoded)) $ \xs -> do
SAT.setVarPolarity solver (xs!0) False
let m = do
ret <- SAT.solve solver
if ret then do
m <- SAT.getModel solver
let sol = decode prob encoded m
SAT.addClause solver $ blockingClause prob encoded sol
return $ Just sol
else
return Nothing
return m
printSolution :: Problem -> Solution -> IO ()
printSolution prob (cells, _) = do
let (w,h,d) = probSize prob
forM_ [1 .. d] $ \z -> do
when (probIsMultiLayer prob) $ do
unless (z == 1) $ putStrLn ""
putStrLn $ "LAYER " ++ show z
forM_ [0 .. h-1] $ \y -> do
putStrLn $ concat $ intersperse ","
[ case Map.lookup (x,y,z) cells of
Nothing -> replicate width '0'
Just k -> replicate (width - length (show k)) '0' ++ show k
| x <- [0 .. w-1]
]
where
width = length $ show (probLineNum prob)
blockingClause :: Problem -> Encoded -> Solution -> SAT.Clause
blockingClause _prob (_,edgesEnc) (_,edges) = [- (edgesEnc Map.! e) | e <- Set.toList edges]
data Options
= Options
{ optOptimize :: Maybe Obj
, optAssumeNoBlank :: Bool
, optAssumeNoDetour :: Bool
, optJSAI2016 :: Bool
, optADC2016 :: Bool
, optNumSolutions :: Int
}
instance Default Options where
def =
Options
{ optOptimize = Nothing
, optAssumeNoBlank = False
, optAssumeNoDetour = False
, optJSAI2016 = False
, optADC2016 = False
, optNumSolutions = 1
}
type Obj = (Bool,Bool)
options :: [OptDescr (Options -> Options)]
options =
[ Option [] ["optimize"]
(OptArg (\arg opt -> opt{ optOptimize = Just (maybe (True,True) readObj arg) }) "<str>")
"Specify objective functions (-n option will be ignored): both (default), length, corners"
, Option [] ["no-blank"]
(NoArg (\opt -> opt{ optAssumeNoBlank = True }))
"Assume no blank cells in solution"
, Option [] ["no-detour"]
(NoArg (\opt -> opt{ optAssumeNoDetour = True }))
"Assume no detour in solution"
, Option [] ["jsai2016"]
(NoArg (\opt -> opt{ optJSAI2016 = True }))
"Add constraints of JSATI2016 paper of Tatsuya Seko et.al."
, Option [] ["adc2016"]
(NoArg (\opt -> opt{ optADC2016 = True }))
"Add constraints specific to ADC2016 rule"
, Option ['n'] []
(ReqArg (\val opt -> opt{ optNumSolutions = read val }) "<int>")
"Maximal number of solutions to enumerate, or negative value for all solutions (default: 1)"
]
where
readObj = f . map toLower
where
f "both" = (True, True)
f "length" = (True, False)
f "corners" = (False, True)
f _ = error "unknown objective"
header :: String
header = unlines
[ "Usage:"
, " numberlink [OPTION]... problem.txt"
, " numberlink [OPTION]... problem.txt encoded.opb"
, ""
, "Options:"
]
showHelp :: Handle -> IO ()
showHelp h = hPutStrLn h (usageInfo header options)
main :: IO ()
main = do
args <- getArgs
case getOpt Permute options args of
(_,_,errs@(_:_)) -> do
mapM_ putStrLn errs
exitFailure
(o,args2,[]) -> do
let opt = foldl (flip id) def o
case args2 of
[fname] -> do
r <- ParsecBL.parseFromFile parser fname
case r of
Left err -> error (show err)
Right prob@Problem{ probSize = (w,h,d) } -> do
putStrLn $ "SIZE " ++ show w ++ "X" ++ show h ++
(if probIsMultiLayer prob then "X" ++ show d else "")
case optOptimize opt of
Nothing -> do
solve <- createSolver opt prob
let loop n
| n == 0 = return ()
| otherwise = do
m <- solve
case m of
Nothing -> return ()
Just sol -> do
printSolution prob sol
putStrLn ""
loop (if n > 0 then n - 1 else n)
loop (optNumSolutions opt)
Just _ -> do
solver <- SAT.newSolver
SAT.setLogger solver $ hPutStrLn stderr
encoded@(cells,_edges) <- encode solver opt prob
unless (optAssumeNoBlank opt) $ do
forM_ (elems cells) $ \xs -> do
SAT.setVarPolarity solver (xs!0) False
obj <- encodeObj solver opt prob encoded
pbo <- PBO.newOptimizer2 solver obj (\m -> evalObj opt prob (decode prob encoded m))
PBO.setLogger pbo $ hPutStrLn stderr
PBO.setOnUpdateBestSolution pbo $ \m val -> do
hPutStrLn stderr $ "# obj = " ++ show val
hFlush stderr
let sol = decode prob encoded m
printSolution prob sol
putStrLn ""
hFlush stdout
PBO.optimize pbo
[fname, fname2] -> do
r <- ParsecBL.parseFromFile parser fname
case r of
Left err -> error (show err)
Right prob -> do
store <- PBStore.newPBStore
obj <- do
encoded <- encode store opt prob
case optOptimize opt of
Nothing -> return Nothing
Just _ -> do
obj <- encodeObj store opt prob encoded
return $ Just [(c,[v]) | (c,v) <- obj]
opb <- PBStore.getPBFormula store
PB.writeOPBFile fname2 $ opb{ PB.pbObjectiveFunction = obj }
_ -> do
showHelp stderr
sampleFile :: BL.ByteString
sampleFile = BL.unlines
[ "SIZE 15X15"
, "LINE_NUM 12"
, "LINE#1 (10,0)-(4,14)"
, "LINE#2 (9,5)-(9,14)"
, "LINE#3 (2,7)-(4,12)"
, "LINE#4 (2,5)-(7,13)"
, "LINE#5 (2,4)-(4,10)"
, "LINE#6 (2,2)-(5,5)"
, "LINE#7 (6,5)-(5,10)"
, "LINE#8 (5,0)-(7,9)"
, "LINE#9 (10,2)-(12,7)"
, "LINE#10 (7,1)-(12,9)"
, "LINE#11 (9,8)-(12,12)"
, "LINE#12 (8,9)-(12,10)"
]
sample :: Problem
Right sample = parse parser "sample" sampleFile