Last active
January 18, 2022 16:41
-
-
Save lzybkr/124b00903dd01b87de25a850da2f1ce3 to your computer and use it in GitHub Desktop.
Using Z3 w/ PowerShell to solve a KenKen puzzle
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| using namespace Microsoft.Z3 | |
| using namespace System.Collections.Generic | |
| using namespace System.Management.Automation.Language | |
| function GetZ3Ast([System.Management.Automation.Language.Ast]$ast, $ctx) | |
| { | |
| $typeName = $ast.GetType().Name | |
| switch ($typeName) | |
| { | |
| "ScriptBlockAst" { | |
| # Ignore some stuff I shouldn't ignore | |
| $t = $ctx.MkTrue() | |
| foreach ($a in $ast.EndBlock.Statements) { | |
| $t = $ctx.MkAnd((GetZ3Ast $a $ctx), $t) | |
| } | |
| return $t | |
| } | |
| "PipelineAst" { | |
| $expr = $ast.GetPureExpression() | |
| if (!$expr) { throw "Only pure expressions supported" } | |
| return GetZ3Ast $expr $ctx | |
| } | |
| "BinaryExpressionAst" { | |
| $left = GetZ3Ast $ast.Left $ctx | |
| $right = GetZ3Ast $ast.Right $ctx | |
| $method = switch ($ast.Operator) | |
| { | |
| "Ieq" { "MkEq"; break } | |
| "Divide" { "MkDiv"; break } | |
| "Minus" { "MkSub"; break } | |
| "Multiply" { "MkMul"; break } | |
| "Plus" { "MkAdd"; break } | |
| default { throw "Unexpected operator $($ast.Operator)" } | |
| } | |
| return $ctx.$method($left, $right) | |
| } | |
| "ParenExpressionAst" { return GetZ3Ast $ast.Pipeline $ctx } | |
| "ConstantExpressionAst" { | |
| # Only support ints for now | |
| return $ctx.MkInt($ast.Value) | |
| } | |
| "VariableExpressionAst" { | |
| $varName = $ast.VariablePath.UserPath | |
| $r = $free_vars[$varName] | |
| if (!$r) { throw "Missing var $varName" } | |
| return $r | |
| } | |
| } | |
| throw "Not yet supported: $typeName" | |
| } | |
| $d = [Dictionary[string, string]]::new() | |
| $d["model"] = "true" | |
| $ctx = [Context]::new($d) | |
| try { | |
| $free_vars = @{} | |
| $b = [IntExpr[][]]::new(9) | |
| for ($i = 0; $i -lt 9; $i++) | |
| { | |
| $b[$i] = [IntExpr[]]::new(9) | |
| for ($j = 0; $j -lt 9; $j++) | |
| { | |
| $varName = "b$i$j" | |
| $b[$i][$j] = $ctx.MkIntConst($varName) | |
| Set-Variable -Name $varName -Value $b[$i][$j] | |
| $free_vars[$varName] = $b[$i][$j] | |
| } | |
| } | |
| $cells_c = [BoolExpr[][]]::new(9) | |
| for ($i = 0; $i -lt 9; $i++) | |
| { | |
| $cells_c[$i] = [BoolExpr[]]::new(9) | |
| for ($j = 0; $j -lt 9; $j++) | |
| { | |
| $cells_c[$i][$j] = $ctx.MkAnd( | |
| $ctx.MkLe($ctx.MkInt(1), $b[$i][$j]), | |
| $ctx.MkLe($b[$i][$j], $ctx.MkInt(9))) | |
| } | |
| } | |
| $rows_c = [BoolExpr[]]::new(9) | |
| for ($i = 0; $i -lt 9; $i++) | |
| { | |
| $rows_c[$i] = $ctx.MkDistinct($b[$i]) | |
| } | |
| $cols_c = [BoolExpr[]]::new(9) | |
| for ($j = 0; $j -lt 9; $j++) | |
| { | |
| $column = [IntExpr[]]::new(9) | |
| for ($i = 0; $i -lt 9; $i++) | |
| { | |
| $column[$i] = $b[$i][$j] | |
| } | |
| $cols_c[$j] = $ctx.MkDistinct($column) | |
| } | |
| $kenken = $ctx.MkTrue() | |
| foreach ($t in $cells_c) | |
| { | |
| $kenken = $ctx.MkAnd($ctx.MkAnd($t), $kenken) | |
| } | |
| $kenken = $ctx.MkAnd($ctx.MkAnd($rows_c), $kenken) | |
| $kenken = $ctx.MkAnd($ctx.MkAnd($cols_c), $kenken) | |
| $ast = { | |
| (($b00*10 + $b01) / $b10) -eq 21 | |
| (($b02*100 + $b03*10 + $b04) * $b14) -eq 1960 | |
| (($b05*10 + $b06) / $b15) -eq 13 | |
| (($b07*10 + $b08) - $b18) -eq 17 | |
| (($b11*100 + $b12*10 + $b13) / (($b22 * 100) + ($b23 * 10) + $b24)) -eq 5 | |
| (($b16*10 + $b17) * (($b27 * 10) + $b28)) -eq 969 | |
| (($b20*10 + $b21) / $b30) -eq 21 | |
| ($b25 - ($b34*10 + $b35)) -eq 66 | |
| ($b27 - $b37) -eq 1 | |
| (($b31*10 + $b32) + ($b40*10 + $b41)) -eq 63 | |
| ($b33 * ($b42*10 + $b43) * $b53) -eq 342 | |
| (($b37*10 + $b38) / $b47) -eq 9 | |
| (($b44*100 + $b45*10 + $b46) * $b56 * ($b65* 10 + $b66)) -eq 59049 | |
| (($b50*100 + $b51*10 + $b52) - ($b60*100 + $b61*10 + $b62)) -eq 73 | |
| (($b54*10 + $b55) - ($b63*10 + $b64)) -eq 1 | |
| ($b48 + ($b57*10 + $b58) + $b68) -eq 57 | |
| } | |
| $x = GetZ3Ast $ast.Ast $ctx | |
| $solver = $ctx.MkSolver() | |
| $solver.Assert($kenken) | |
| $solver.Assert($x) | |
| if ($solver.Check() -eq 'SATISFIABLE') | |
| { | |
| 'SATISFIABLE' | |
| ($m = $solver.Model) | |
| } | |
| } finally { | |
| $ctx.Dispose() | |
| } |
This is my code to solve this Jane Street puzzle.
The AST isn't strictly necessary, it was just a concise way of expressing the constraints that are easily expressed as math. This is roughly what it might look like without the AST (created by modifying the script slightly to return strings instead of actually making the calls):
$ctx.MkAnd(
$ctx.MkEq($ctx.MkAdd($ctx.MkAdd($cell[4][8], $ctx.MkAdd($ctx.MkMul($cell[5][7], $ctx.MkInt(10)), $cell[5][8])), $cell[6][8]), $ctx.MkInt(57)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkSub($ctx.MkAdd($ctx.MkMul($cell[5][4], $ctx.MkInt(10)), $cell[5][5]), $ctx.MkAdd($ctx.MkMul($cell[6][3], $ctx.MkInt(10)), $cell[6][4])), $ctx.MkInt(1)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkSub($ctx.MkAdd($ctx.MkAdd($ctx.MkMul($cell[5][0], $ctx.MkInt(100)), $ctx.MkMul($cell[5][1], $ctx.MkInt(10))), $cell[5][2]), $ctx.MkAdd($ctx.MkAdd($ctx.MkMul($cell[6][0], $ctx.MkInt(100)), $ctx.MkMul($cell[6][1], $ctx.MkInt(10))), $cell[6][2])), $ctx.MkInt(73)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkMul($ctx.MkMul($ctx.MkAdd($ctx.MkAdd($ctx.MkMul($cell[4][4], $ctx.MkInt(100)), $ctx.MkMul($cell[4][5], $ctx.MkInt(10))), $cell[4][6]), $cell[5][6]), $ctx.MkAdd($ctx.MkMul($cell[6][5], $ctx.MkInt(10)), $cell[6][6])), $ctx.MkInt(59049)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkDiv($ctx.MkAdd($ctx.MkMul($cell[3][7], $ctx.MkInt(10)), $cell[3][8]), $cell[4][7]), $ctx.MkInt(9)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkMul($ctx.MkMul($cell[3][3], $ctx.MkAdd($ctx.MkMul($cell[4][2], $ctx.MkInt(10)), $cell[4][3])), $cell[5][3]), $ctx.MkInt(342)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkAdd($ctx.MkAdd($ctx.MkMul($cell[3][1], $ctx.MkInt(10)), $cell[3][2]), $ctx.MkAdd($ctx.MkMul($cell[4][0], $ctx.MkInt(10)), $cell[4][1])), $ctx.MkInt(63)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkSub($cell[2][7], $cell[3][7]), $ctx.MkInt(1)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkSub($cell[2][5], $ctx.MkAdd($ctx.MkMul($cell[3][4], $ctx.MkInt(10)), $cell[3][5])), $ctx.MkInt(66)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkDiv($ctx.MkAdd($ctx.MkMul($cell[2][0], $ctx.MkInt(10)), $cell[2][1]), $cell[3][0]), $ctx.MkInt(21)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkMul($ctx.MkAdd($ctx.MkMul($cell[1][6], $ctx.MkInt(10)), $cell[1][7]), $ctx.MkAdd($ctx.MkMul($cell[2][7], $ctx.MkInt(10)), $cell[2][8])), $ctx.MkInt(969)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkDiv($ctx.MkAdd($ctx.MkAdd($ctx.MkMul($cell[1][1], $ctx.MkInt(100)), $ctx.MkMul($cell[1][2], $ctx.MkInt(10))), $cell[1][3]), $ctx.MkAdd($ctx.MkAdd($ctx.MkMul($cell[2][2], $ctx.MkInt(100)), $ctx.MkMul($cell[2][3], $ctx.MkInt(10))), $cell[2][4])), $ctx.MkInt(5)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkSub($ctx.MkAdd($ctx.MkMul($cell[0][7], $ctx.MkInt(10)), $cell[0][8]), $cell[1][8]), $ctx.MkInt(17)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkDiv($ctx.MkAdd($ctx.MkMul($cell[0][5], $ctx.MkInt(10)), $cell[0][6]), $cell[1][5]), $ctx.MkInt(13)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkMul($ctx.MkAdd($ctx.MkAdd($ctx.MkMul($cell[0][2], $ctx.MkInt(100)), $ctx.MkMul($cell[0][3], $ctx.MkInt(10))), $cell[0][4]), $cell[1][4]), $ctx.MkInt(1960)),
$ctx.MkAnd(
$ctx.MkEq($ctx.MkDiv($ctx.MkAdd($ctx.MkMul($cell[0][0], $ctx.MkInt(10)), $cell[0][1]), $cell[1][0]), $ctx.MkInt(21)),
$ctx.MkTrue()))))))))))))))))
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Hello Jason
oh, this is cool, can you, please, explain how it works and why AST is here? :)
I'm trying to see Z3 from PowerShell for a different purpose, and so fat all that z3 stuff looks pretty cryptic :)