1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3#
4# Copyright (C) 2024 Antmicro
5#
6# Licensed under the Apache License, Version 2.0 (the "License");
7# you may not use this file except in compliance with the License.
8# You may obtain a copy of the License at
9#
10#     http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
17#
18# SPDX-License-Identifier: Apache-2.0
19
20from .ast import *
21from . import operators as op
22from .helper import Visitor
23
24class EvalNulls(Visitor):
25    def __init__(self, nodes: Node, verbose: bool = False) -> None:
26        super().__init__(nodes, verbose)
27
28    def visit_OR(self, node: op.OR) -> None:
29        self.iterate_children_dfs(node)
30
31        match (node.lhs, node.rhs):
32            case (IntLit() as lit, expr) | (expr, IntLit() as lit):
33                if lit.value == 0:
34                    node.replace(expr.detach())
35            case (IntLit() as lit1, IntLit() as lit2):
36                if lit1.value == 0 and lit2.value == 0:
37                    node.replace(lit1.detach())
38
39    def visit_AND(self, node: op.AND) -> None:
40        self.iterate_children_dfs(node)
41
42        match (node.lhs, node.rhs):
43            case (IntLit() as lit, expr) | (expr, IntLit() as lit):
44                if lit.value == 0:
45                    node.replace(lit.detach())
46                elif lit.value == (1 << expr.type.width) - 1:
47                    node.replace(expr.detach())
48
49    def visit_BinaryOp(self, node: BinaryOp) -> None:
50        self.iterate_children_dfs(node)
51
52        match node:
53            case op.SHL() | op.SHR() | op.USHR():
54                if isinstance(node.rhs, IntLit) and node.rhs.value == 0:
55                    node.replace(node.lhs.detach())
56            case op.Add() | op.Sub():
57                match (node.lhs, node.rhs):
58                    case (IntLit() as lit, other) | (other, IntLit() as lit):
59                        if lit.value == 0:
60                            node.replace(lit.detach())
61