/** Utility definitions for reasoning about equality methods. */

import python
import semmle.python.dataflow.new.DataFlow

private Attribute dictAccess(LocalVariable var) {
  result.getName() = "__dict__" and
  result.getObject() = var.getAnAccess()
}

private Call getattr(LocalVariable obj, LocalVariable attr) {
  result.getFunc().(Name).getId() = "getattr" and
  result.getArg(0) = obj.getAnAccess() and
  result.getArg(1) = attr.getAnAccess()
}

/**
 * A generic equality method that compares all attributes in its dict,
 * or compares attributes using `getattr`.
 */
class GenericEqMethod extends Function {
  GenericEqMethod() {
    this.getName() = "__eq__" and
    exists(LocalVariable self, LocalVariable other |
      self.getAnAccess() = this.getArg(0) and
      self.getId() = "self" and
      other.getAnAccess() = this.getArg(1) and
      exists(Compare eq |
        eq.getOp(0) instanceof Eq or
        eq.getOp(0) instanceof NotEq
      |
        // `self.__dict__ == other.__dict__`
        eq.getAChildNode() = dictAccess(self) and
        eq.getAChildNode() = dictAccess(other)
        or
        // `getattr(self, var) == getattr(other, var)`
        exists(Variable var |
          eq.getAChildNode() = getattr(self, var) and
          eq.getAChildNode() = getattr(other, var)
        )
      )
    )
  }
}

/** An `__eq__` method that just does `self is other` */
class IdentityEqMethod extends Function {
  IdentityEqMethod() {
    this.getName() = "__eq__" and
    exists(LocalVariable self, LocalVariable other |
      self.getAnAccess() = this.getArg(0) and
      self.getId() = "self" and
      other.getAnAccess() = this.getArg(1) and
      exists(Compare eq | eq.getOp(0) instanceof Is |
        eq.getAChildNode() = self.getAnAccess() and
        eq.getAChildNode() = other.getAnAccess()
      )
    )
  }
}

/** An (in)equality method that delegates to its complement */
class DelegatingEqualityMethod extends Function {
  DelegatingEqualityMethod() {
    exists(Return ret, UnaryExpr not_, Expr comp, Parameter p0, Parameter p1 |
      ret.getScope() = this and
      ret.getValue() = not_ and
      not_.getOp() instanceof Not and
      not_.getOperand() = comp
    |
      exists(Cmpop op |
        comp.(Compare).compares(p0.getVariable().getAnAccess(), op, p1.getVariable().getAnAccess())
      |
        this.getName() = "__eq__" and op instanceof NotEq
        or
        this.getName() = "__ne__" and op instanceof Eq
      )
      or
      exists(DataFlow::MethodCallNode call, string name |
        call.calls(DataFlow::exprNode(p0.getVariable().getAnAccess()), name) and
        call.getArg(0).asExpr() = p1.getVariable().getAnAccess()
      |
        this.getName() = "__eq__" and name = "__ne__"
        or
        this.getName() = "__ne__" and name = "__eq__"
      )
    )
  }
}