import python

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_, Compare comp, Cmpop op, Parameter p0, Parameter p1 |
      ret.getScope() = this and
      ret.getValue() = not_ and
      not_.getOp() instanceof Not and
      not_.getOperand() = comp and
      comp.compares(p0.getVariable().getAnAccess(), op, p1.getVariable().getAnAccess())
    |
      this.getName() = "__eq__" and op instanceof NotEq
      or
      this.getName() = "__ne__" and op instanceof Eq
    )
  }
}