Coverage for src / hodoku / core / cell_set.py: 100%

1"""81-cell bitset backed by a Python int.""" 

2 

3from __future__ import annotations 

4 

5from typing import Iterator 

6 

7_ALL_BITS = (1 << 81) - 1 

8 

9 

10class CellSet: 

11 """Bitset over the 81 cells of a sudoku grid (indices 0-80). 

12 

13 Backed by a single Python int. Bit i corresponds to cell i. 

14 """ 

15 

16 __slots__ = ("_bits",) 

17 

18 def __init__(self, bits: int = 0) -> None: 

19 self._bits = bits 

20 

21 # --- mutation --- 

22 

23 def add(self, index: int) -> None: 

24 self._bits |= 1 << index 

25 

26 def remove(self, index: int) -> None: 

27 self._bits &= ~(1 << index) 

28 

29 def clear(self) -> None: 

30 self._bits = 0 

31 

32 def set_all(self) -> None: 

33 self._bits = _ALL_BITS 

34 

35 def set(self, other: CellSet) -> None: 

36 self._bits = other._bits 

37 

38 # --- set operations (in-place) --- 

39 

40 def and_(self, other: CellSet) -> None: 

41 self._bits &= other._bits 

42 

43 def or_(self, other: CellSet) -> None: 

44 self._bits |= other._bits 

45 

46 def and_not(self, other: CellSet) -> None: 

47 self._bits &= ~other._bits 

48 

49 # --- set operations (returning new CellSet) --- 

50 

51 def intersection(self, other: CellSet) -> CellSet: 

52 return CellSet(self._bits & other._bits) 

53 

54 def union(self, other: CellSet) -> CellSet: 

55 return CellSet(self._bits | other._bits) 

56 

57 def difference(self, other: CellSet) -> CellSet: 

58 return CellSet(self._bits & ~other._bits) 

59 

60 # --- query --- 

61 

62 def contains(self, index: int) -> bool: 

63 return bool(self._bits >> index & 1) 

64 

65 def is_empty(self) -> bool: 

66 return self._bits == 0 

67 

68 def size(self) -> int: 

69 return self._bits.bit_count() 

70 

71 def get(self, n: int) -> int: 

72 """Return the index of the nth set bit (0-based).""" 

73 bits = self._bits 

74 for _ in range(n): 

75 bits &= bits - 1 # clear lowest set bit 

76 return (bits & -bits).bit_length() - 1 

77 

78 def first(self) -> int: 

79 """Return the lowest set index, or -1 if empty.""" 

80 if self._bits == 0: 

81 return -1 

82 return (self._bits & -self._bits).bit_length() - 1 

83 

84 def equals(self, other: CellSet) -> bool: 

85 return self._bits == other._bits 

86 

87 # --- iteration --- 

88 

89 def __iter__(self) -> Iterator[int]: 

90 bits = self._bits 

91 while bits: 

92 lsb = bits & -bits 

93 yield lsb.bit_length() - 1 

94 bits ^= lsb 

95 

96 def __len__(self) -> int: 

97 return self._bits.bit_count() 

98 

99 def __eq__(self, other: object) -> bool: 

100 if isinstance(other, CellSet): 

101 return self._bits == other._bits 

102 return NotImplemented 

103 

104 def __hash__(self) -> int: 

105 return hash(self._bits) 

106 

107 def __repr__(self) -> str: 

108 indices = list(self) 

109 return f"CellSet({indices})" 

110 

111 def clone(self) -> CellSet: 

112 return CellSet(self._bits)